├── README.md ├── level1 └── train.py └── level2 └── train_customData.py /README.md: -------------------------------------------------------------------------------- 1 | ### There are some examples of using PyTorch for image classification 2 | 3 | # Usage 4 | 5 | **Each file of this project is an example of image classification, you can learn from level1 to levelN. For more explaination of these codes, please visit [CSDN Blog](http://blog.csdn.net/u014380165/article/category/7286599)** 6 | 7 | * level1: Using default Dataset class(`torchvision.datasets.ImageFolder`) to read Image. This code is modified from [PyTorch](http://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html) 8 | 9 | * level2: Using custom Dataset class(a custom class inherit from base class `torch.utils.data.Dataset`) to read Image. -------------------------------------------------------------------------------- /level1/train.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, division 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.optim as optim 6 | from torch.optim import lr_scheduler 7 | from torch.autograd import Variable 8 | import torchvision 9 | from torchvision import datasets, models, transforms 10 | import time 11 | import os 12 | 13 | def train_model(model, criterion, optimizer, scheduler, num_epochs=25): 14 | since = time.time() 15 | 16 | best_model_wts = model.state_dict() 17 | best_acc = 0.0 18 | 19 | for epoch in range(num_epochs): 20 | print('Epoch {}/{}'.format(epoch, num_epochs - 1)) 21 | print('-' * 10) 22 | 23 | # Each epoch has a training and validation phase 24 | for phase in ['train', 'val']: 25 | if phase == 'train': 26 | scheduler.step() 27 | model.train(True) # Set model to training mode 28 | else: 29 | model.train(False) # Set model to evaluate mode 30 | 31 | running_loss = 0.0 32 | running_corrects = 0.0 33 | 34 | # Iterate over data. 35 | for data in dataloders[phase]: 36 | # get the inputs 37 | inputs, labels = data 38 | 39 | # wrap them in Variable 40 | if use_gpu: 41 | inputs = Variable(inputs.cuda()) 42 | labels = Variable(labels.cuda()) 43 | else: 44 | inputs, labels = Variable(inputs), Variable(labels) 45 | 46 | # zero the parameter gradients 47 | optimizer.zero_grad() 48 | 49 | # forward 50 | outputs = model(inputs) 51 | _, preds = torch.max(outputs.data, 1) 52 | loss = criterion(outputs, labels) 53 | 54 | # backward + optimize only if in training phase 55 | if phase == 'train': 56 | loss.backward() 57 | optimizer.step() 58 | 59 | # statistics 60 | running_loss += loss.data[0] 61 | running_corrects += torch.sum(preds == labels.data).to(torch.float32) 62 | 63 | epoch_loss = running_loss / dataset_sizes[phase] 64 | epoch_acc = running_corrects / dataset_sizes[phase] 65 | 66 | print('{} Loss: {:.4f} Acc: {:.4f}'.format( 67 | phase, epoch_loss, epoch_acc)) 68 | 69 | # deep copy the model 70 | if phase == 'val' and epoch_acc > best_acc: 71 | best_acc = epoch_acc 72 | best_model_wts = model.state_dict() 73 | 74 | time_elapsed = time.time() - since 75 | print('Training complete in {:.0f}m {:.0f}s'.format( 76 | time_elapsed // 60, time_elapsed % 60)) 77 | print('Best val Acc: {:4f}'.format(best_acc)) 78 | 79 | # load best model weights 80 | model.load_state_dict(best_model_wts) 81 | return model 82 | 83 | if __name__ == '__main__': 84 | 85 | # data_transform, pay attention that the input of Normalize() is Tensor and the input of RandomResizedCrop() or RandomHorizontalFlip() is PIL Image 86 | data_transforms = { 87 | 'train': transforms.Compose([ 88 | transforms.RandomSizedCrop(224), 89 | transforms.RandomHorizontalFlip(), 90 | transforms.ToTensor(), 91 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 92 | ]), 93 | 'val': transforms.Compose([ 94 | transforms.Scale(256), 95 | transforms.CenterCrop(224), 96 | transforms.ToTensor(), 97 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 98 | ]), 99 | } 100 | 101 | # your image data file 102 | data_dir = '/data' 103 | image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), 104 | data_transforms[x]) for x in ['train', 'val']} 105 | # wrap your data and label into Tensor 106 | dataloders = {x: torch.utils.data.DataLoader(image_datasets[x], 107 | batch_size=4, 108 | shuffle=True, 109 | num_workers=4) for x in ['train', 'val']} 110 | 111 | dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} 112 | 113 | # use gpu or not 114 | use_gpu = torch.cuda.is_available() 115 | 116 | # get model and replace the original fc layer with your fc layer 117 | model_ft = models.resnet18(pretrained=True) 118 | num_ftrs = model_ft.fc.in_features 119 | model_ft.fc = nn.Linear(num_ftrs, 2) 120 | 121 | if use_gpu: 122 | model_ft = model_ft.cuda() 123 | 124 | # define loss function 125 | criterion = nn.CrossEntropyLoss() 126 | 127 | # Observe that all parameters are being optimized 128 | optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9) 129 | 130 | # Decay LR by a factor of 0.1 every 7 epochs 131 | exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1) 132 | 133 | model_ft = train_model(model=model_ft, 134 | criterion=criterion, 135 | optimizer=optimizer_ft, 136 | scheduler=exp_lr_scheduler, 137 | num_epochs=25) 138 | -------------------------------------------------------------------------------- /level2/train_customData.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, division 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.optim as optim 6 | from torch.optim import lr_scheduler 7 | from torch.autograd import Variable 8 | from torchvision import models, transforms 9 | 10 | import time 11 | import os 12 | from torch.utils.data import Dataset 13 | 14 | from PIL import Image 15 | 16 | # use PIL Image to read image 17 | def default_loader(path): 18 | try: 19 | img = Image.open(path) 20 | return img.convert('RGB') 21 | except: 22 | print("Cannot read image: {}".format(path)) 23 | 24 | # define your Dataset. Assume each line in your .txt file is [name/tab/label], for example:0001.jpg 1 25 | class customData(Dataset): 26 | def __init__(self, img_path, txt_path, dataset = '', data_transforms=None, loader = default_loader): 27 | with open(txt_path) as input_file: 28 | lines = input_file.readlines() 29 | self.img_name = [os.path.join(img_path, line.strip().split('\t')[0]) for line in lines] 30 | self.img_label = [int(line.strip().split('\t')[-1]) for line in lines] 31 | self.data_transforms = data_transforms 32 | self.dataset = dataset 33 | self.loader = loader 34 | 35 | def __len__(self): 36 | return len(self.img_name) 37 | 38 | def __getitem__(self, item): 39 | img_name = self.img_name[item] 40 | label = self.img_label[item] 41 | img = self.loader(img_name) 42 | 43 | if self.data_transforms is not None: 44 | try: 45 | img = self.data_transforms[self.dataset](img) 46 | except: 47 | print("Cannot transform image: {}".format(img_name)) 48 | return img, label 49 | 50 | def train_model(model, criterion, optimizer, scheduler, num_epochs, use_gpu): 51 | since = time.time() 52 | 53 | best_model_wts = model.state_dict() 54 | best_acc = 0.0 55 | 56 | for epoch in range(num_epochs): 57 | begin_time = time.time() 58 | print('Epoch {}/{}'.format(epoch, num_epochs - 1)) 59 | print('-' * 10) 60 | 61 | # Each epoch has a training and validation phase 62 | for phase in ['train', 'val']: 63 | count_batch = 0 64 | if phase == 'train': 65 | scheduler.step() 66 | model.train(True) # Set model to training mode 67 | else: 68 | model.train(False) # Set model to evaluate mode 69 | 70 | running_loss = 0.0 71 | running_corrects = 0.0 72 | 73 | # Iterate over data. 74 | for data in dataloders[phase]: 75 | count_batch += 1 76 | # get the inputs 77 | inputs, labels = data 78 | 79 | # wrap them in Variable 80 | if use_gpu: 81 | inputs = Variable(inputs.cuda()) 82 | labels = Variable(labels.cuda()) 83 | else: 84 | inputs, labels = Variable(inputs), Variable(labels) 85 | 86 | # zero the parameter gradients 87 | optimizer.zero_grad() 88 | 89 | # forward 90 | outputs = model(inputs) 91 | _, preds = torch.max(outputs.data, 1) 92 | loss = criterion(outputs, labels) 93 | 94 | # backward + optimize only if in training phase 95 | if phase == 'train': 96 | loss.backward() 97 | optimizer.step() 98 | # statistics 99 | running_loss += loss.data[0] 100 | running_corrects += torch.sum(preds == labels.data).to(torch.float32) 101 | 102 | # print result every 10 batch 103 | if count_batch%10 == 0: 104 | batch_loss = running_loss / (batch_size*count_batch) 105 | batch_acc = running_corrects / (batch_size*count_batch) 106 | print('{} Epoch [{}] Batch [{}] Loss: {:.4f} Acc: {:.4f} Time: {:.4f}s'. \ 107 | format(phase, epoch, count_batch, batch_loss, batch_acc, time.time()-begin_time)) 108 | begin_time = time.time() 109 | 110 | epoch_loss = running_loss / dataset_sizes[phase] 111 | epoch_acc = running_corrects / dataset_sizes[phase] 112 | print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc)) 113 | 114 | # save model 115 | if phase == 'train': 116 | if not os.path.exists('output'): 117 | os.makedirs('output') 118 | torch.save(model, 'output/resnet_epoch{}.pkl'.format(epoch)) 119 | 120 | # deep copy the model 121 | if phase == 'val' and epoch_acc > best_acc: 122 | best_acc = epoch_acc 123 | best_model_wts = model.state_dict() 124 | 125 | time_elapsed = time.time() - since 126 | print('Training complete in {:.0f}m {:.0f}s'.format( 127 | time_elapsed // 60, time_elapsed % 60)) 128 | print('Best val Acc: {:4f}'.format(best_acc)) 129 | 130 | # load best model weights 131 | model.load_state_dict(best_model_wts) 132 | return model 133 | 134 | if __name__ == '__main__': 135 | 136 | data_transforms = { 137 | 'train': transforms.Compose([ 138 | transforms.RandomSizedCrop(224), 139 | transforms.RandomHorizontalFlip(), 140 | transforms.ToTensor(), 141 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 142 | ]), 143 | 'val': transforms.Compose([ 144 | transforms.Scale(256), 145 | transforms.CenterCrop(224), 146 | transforms.ToTensor(), 147 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 148 | ]), 149 | } 150 | 151 | use_gpu = torch.cuda.is_available() 152 | 153 | batch_size = 32 154 | num_class = 2 155 | 156 | image_datasets = {x: customData(img_path='/ImagePath', 157 | txt_path=('/TxtFile/' + x + '.txt'), 158 | data_transforms=data_transforms, 159 | dataset=x) for x in ['train', 'val']} 160 | 161 | # wrap your data and label into Tensor 162 | dataloders = {x: torch.utils.data.DataLoader(image_datasets[x], 163 | batch_size=batch_size, 164 | shuffle=True) for x in ['train', 'val']} 165 | 166 | dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} 167 | 168 | # get model and replace the original fc layer with your fc layer 169 | model_ft = models.resnet50(pretrained=True) 170 | num_ftrs = model_ft.fc.in_features 171 | model_ft.fc = nn.Linear(num_ftrs, num_class) 172 | 173 | # if use gpu 174 | if use_gpu: 175 | model_ft = model_ft.cuda() 176 | 177 | # define cost function 178 | criterion = nn.CrossEntropyLoss() 179 | 180 | # Observe that all parameters are being optimized 181 | optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.005, momentum=0.9) 182 | 183 | # Decay LR by a factor of 0.2 every 5 epochs 184 | exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.2) 185 | 186 | # multi-GPU 187 | model_ft = torch.nn.DataParallel(model_ft, device_ids=[0,1]) 188 | 189 | # train model 190 | model_ft = train_model(model=model_ft, 191 | criterion=criterion, 192 | optimizer=optimizer_ft, 193 | scheduler=exp_lr_scheduler, 194 | num_epochs=25, 195 | use_gpu=use_gpu) 196 | 197 | # save best model 198 | torch.save(model_ft,"output/best_resnet.pkl") 199 | --------------------------------------------------------------------------------