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 # 定义超参 batch_size = 16 num_epochs = 300 # 因为要让他过拟合,epoch数要多一些 # 加载数据 dataset = 'Palmprint' train_directory = os.path.join(dataset, 'train') valid_directory = os.path.join(dataset, 'valid') # 直接跑,不用数据增强,在验证集上75%的准确率,数据增强全部打开,84%的准确率 image_transforms = { 'train': transforms.Compose([ transforms.Grayscale(num_output_channels=3), transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)), transforms.RandomRotation(degrees=15), transforms.CenterCrop(size=224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'valid': transforms.Compose([ transforms.Grayscale(num_output_channels=3), transforms.Resize(size=256), transforms.CenterCrop(size=224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } data = { 'train': datasets.ImageFolder(root=train_directory, transform=image_transforms['train']), 'valid': datasets.ImageFolder(root=valid_directory, transform=image_transforms['valid']) } train_data_size = len(data['train']) valid_data_size = len(data['valid']) train_data = DataLoader(data['train'], batch_size=batch_size, shuffle=True) valid_data = DataLoader(data['valid'], batch_size=batch_size, shuffle=True) print(train_data_size, valid_data_size) # 迁移学习 resnet18 = models.resnet18(pretrained=True) for param in resnet18.parameters(): param.requires_grad = False fc_inputs = resnet18.fc.in_features resnet18.fc = nn.Sequential( nn.Linear(fc_inputs, 256), nn.ReLU(), nn.Dropout(0.4), nn.Linear(256, 99), nn.LogSoftmax(dim=1) ) # 用GPU进行训练 resnet18 = resnet18.to('cuda:0') # 定义损失函数和优化器 loss_func = nn.NLLLoss() optimizer = optim.Adam(resnet18.parameters()) # 训练和验证 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 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_valid_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)) model_path = 'models_resnet18_ep' + str(num_epochs) if not os.path.exists(model_path): os.makedirs(model_path) torch.save(model, model_path + '/' + dataset + '_model_' + str(epoch + 1) + '.pt') return model, history, best_acc, best_epoch trained_model, history, best_acc, best_epoch = train_and_valid(resnet18, loss_func, optimizer, num_epochs) model_path = 'models_resnet18_ep' + str(num_epochs) torch.save(history, model_path + '/' + 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, 8) plt.title('best_acc:' + str(best_acc)[0:6] + ' best_epoch' + str(best_epoch)) plt.savefig(dataset + model_path + '_loss_curve.png') plt.close() plt.plot(history[:, 2:4]) plt.legend(['Tr Accuracy', 'Val Accuracy']) plt.xlabel('Epoch Number') plt.ylabel('Accuracy') plt.ylim(0, 1) plt.title('best_acc:' + str(best_acc)[0:6] + ' best_epoch' + str(best_epoch)) plt.savefig(dataset + model_path + '_accuracy_curve.png') plt.close() # plt.show()
利用RESET实现掌纹识别
猜你喜欢
转载自blog.csdn.net/m0_67250370/article/details/135088675
今日推荐
周排行