暂时不解释,仅作为记录:
import torch from torchvision import datasets, models, transforms import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader import time import numpy as np import matplotlib.pyplot as plt import os image_transforms = { 'cut_image_classify': transforms.Compose([ transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)), transforms.RandomRotation(degrees=15), transforms.RandomHorizontalFlip(), transforms.CenterCrop(size=224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'cut_image_test': transforms.Compose([ transforms.Resize(size=256), transforms.CenterCrop(size=224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } dataset = 'datasets' # train_directory = os.path.join(dataset, 'train') train_directory = os.path.join(dataset, 'cut_image_classify') # valid_directory = os.path.join(dataset, 'valid') valid_directory = os.path.join(dataset, 'cut_image_test') batch_size = 4 num_classes = 5 data = { 'cut_image_classify': datasets.ImageFolder(root=train_directory, transform=image_transforms['cut_image_classify']), 'cut_image_test': datasets.ImageFolder(root=valid_directory, transform=image_transforms['cut_image_test']) } train_data_size = len(data['cut_image_classify']) valid_data_size = len(data['cut_image_test']) train_data = DataLoader(data['cut_image_classify'], batch_size=batch_size, shuffle=True) valid_data = DataLoader(data['cut_image_test'], batch_size=batch_size, shuffle=True) print(train_data_size, valid_data_size) ResNet50 = models.resnet50(pretrained=True) # 在PyTorch中加载模型时,所有参数的‘requires_grad’字段默认设置为true。 # 这意味着对参数值的每一次更改都将被存储,以便在用于训练的反向传播图中使用。 # 这增加了内存需求。 # 由于预训练的模型中的大多数参数已经训练好了,因此将requires_grad字段重置为false。 for param in ResNet50.parameters(): param.requires_grad = False fc_inputs = ResNet50.fc.in_features ResNet50.fc = nn.Sequential( nn.Linear(fc_inputs, 256), nn.ReLU(), nn.Dropout(0.4), nn.Linear(256, 5), nn.LogSoftmax(dim=1) ) # # 用GPU进行训练 # resnet50 = resnet50.to('cuda:0') # 用CPU训练 resnet50 = ResNet50.to('cpu') # loss_func = nn.CrossEntropyLoss() loss_func = nn.NLLLoss() optimizer = optim.Adam(resnet50.parameters(), lr=0.0002) def train_and_valid(model, loss_function, optimizer, epochs=25): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") history = [] best_acc = 0.0 best_epoch = 0 for epoch in range(epochs): epoch_start = time.time() print("Epoch: {}/{}".format(epoch + 1, epochs)) model.train() train_loss = 0.0 train_acc = 0.0 valid_loss = 0.0 valid_acc = 0.0 for i, (inputs, labels) in enumerate(train_data): inputs = inputs.to(device) labels = labels.to(device) # 因为这里梯度是累加的,所以每次记得清零 optimizer.zero_grad() outputs = model(inputs) loss = loss_function(outputs, labels) loss.backward() optimizer.step() train_loss += loss.item() * inputs.size(0) ret, predictions = torch.max(outputs.data, 1) correct_counts = predictions.eq(labels.data.view_as(predictions)) acc = torch.mean(correct_counts.type(torch.FloatTensor)) train_acc += acc.item() * inputs.size(0) with torch.no_grad(): model.eval() for j, (inputs, labels) in enumerate(valid_data): inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) loss = loss_function(outputs, labels) valid_loss += loss.item() * inputs.size(0) ret, predictions = torch.max(outputs.data, 1) correct_counts = predictions.eq(labels.data.view_as(predictions)) acc = torch.mean(correct_counts.type(torch.FloatTensor)) valid_acc += acc.item() * inputs.size(0) avg_train_loss = train_loss / train_data_size avg_train_acc = train_acc / train_data_size avg_valid_loss = valid_loss / valid_data_size avg_valid_acc = valid_acc / valid_data_size history.append([avg_train_loss, avg_valid_loss, avg_train_acc, avg_valid_acc]) if best_acc < avg_valid_acc: best_acc = avg_valid_acc best_epoch = epoch + 1 torch.save(model, 'models/' + dataset + '_model_' + str(epoch + 1) + '.pt') epoch_end = time.time() print( "Epoch: {:03d}, Training: Loss: {:.4f}, Accuracy: {:.4f}%, \n\t\tValidation: Loss: {:.4f}, Accuracy: {:.4f}%, Time: {:.4f}s".format( epoch + 1, avg_train_loss, avg_train_acc * 100, avg_valid_loss, avg_valid_acc * 100, epoch_end - epoch_start )) print("Best Accuracy for validation : {:.4f} at epoch {:03d}".format(best_acc, best_epoch)) # torch.save(model, 'models/' + dataset + '_model_' + str(epoch + 1) + '.pt') return model, history if __name__ == "__main__": num_epochs = 600 trained_model, history = train_and_valid(resnet50, loss_func, optimizer, num_epochs) torch.save(history, 'models/' + dataset + '_history.pt') history = np.array(history) plt.plot(history[:, 0:2]) plt.legend(['Tr Loss', 'Val Loss']) plt.xlabel('Epoch Number') plt.ylabel('Loss') plt.ylim(0, 2) plt.savefig(dataset + '_loss_curve.png') plt.show() plt.plot(history[:, 2:4]) plt.legend(['Tr Accuracy', 'Val Accuracy']) plt.xlabel('Epoch Number') plt.ylabel('Accuracy') plt.ylim(0, 1) plt.savefig(dataset + '_accuracy_curve.png') plt.show()