dataset
from PIL import Image from torch.utils.data import Dataset import os class MyDataLoader(Dataset): def __init__(self, root_dir, transform, augment=True): self.root_dir = root_dir self.augment = augment self.testNameToNum = {"ok": 0, "wrong": 1} self.numToName = {0: "良品", 1: "连锡"} self.trainNameToNum = {"良品": 0, "连锡": 1} # 判断文件夹dirname是否存在 if not os.path.exists(self.root_dir): print("error: folder \"", self.root_dir, "\" not exits!") raise Exception("error: folder" + self.root_dir + "文件夹" + "not exits!") self.names = os.listdir(self.root_dir) # 创建5*2的数据集 self.images_path = self.get_all_images_path() # 5个数据的标签 self.images_labels = self.get_all_images_labels() self.images_path = sum(self.images_path, []) self.transform = transform def get_all_images_path(self): all_images_path = [] for root, dirs, files in os.walk(self.root_dir): every_dir_images_path = [] for file in files: path = os.path.join(root, file) every_dir_images_path.append(path) all_images_path.append(every_dir_images_path) return all_images_path[1:] def get_all_images_labels(self): all_images_labels = [] for i in range(len(self.names)): every_dir_images_path = self.images_path[i] every_dir_images_labels = [] if self.names[i] in self.testNameToNum: every_dir_images_labels = [self.testNameToNum[self.names[i]]] * len(every_dir_images_path) elif self.names[i] in self.trainNameToNum: every_dir_images_labels = [self.trainNameToNum[self.names[i]]] * len(every_dir_images_path) if every_dir_images_labels is None: print("error: folder \"", self.root_dir, "\" not exits!") raise Exception(self.names[i] + "没有符合的标签") all_images_labels.append(every_dir_images_labels) return sum(all_images_labels, []) # 根据索引获取data和label def __getitem__(self, index): ig_path = self.images_path[index] label = self.images_labels[index] ###################################### pil_image = Image.open(ig_path).convert('RGB') ######################################### data = self.transform(pil_image) return data, label # 获取数据集的大小 def __len__(self): return len(self.images_path) # if __name__ == '__main__': # dirname = "data//test" # data = MyDataLoader(dirname) # print(f'data size is : {len(data)}') # # image,label = data[1] # print(str(data))
train
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 from tqdm import tqdm from MyDataLoader import MyDataLoader transform = 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]) ]) # 三、加载数据 # torchvision.transforms包DataLoader是 Pytorch 重要的特性，它们使得数据增加和加载数据变得非常简单。 # 使用 DataLoader 加载数据的时候就会将之前定义的数据 transform 就会应用的数据上了。 dataset = 'data' train_directory = os.path.join(dataset, 'train') test_directory = os.path.join(dataset, 'test') train_dataset = MyDataLoader(train_directory, transform) test_dataset = MyDataLoader(test_directory, transform) train_data_size = len(train_dataset) test_data_size = len(test_dataset) batch_size = 64 num_classes = 6 train_data = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_data = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) # 四、迁移学习 # 这里使用ResNet-50的预训练模型。 resnet50 = models.resnet50(pretrained=True) # 在PyTorch中加载模型时，所有参数的‘requires_grad’字段默认设置为true。 # 这意味着对参数值的每一次更改都将被存储，以便在用于训练的反向传播图中使用。 # 这增加了内存需求。由于预训练的模型中的大多数参数已经训练好了，因此将requires_grad字段重置为false。 for param in resnet50.parameters(): param.requires_grad = False # 为了适应自己的数据集，将ResNet-50的最后一层替换为， # 将原来最后一个全连接层的输入喂给一个有256个输出单元的线性层，接着再连接ReLU层和Dropout层， # 然后是256 x 6的线性层，输出为6通道的softmax层。 fc_inputs = resnet50.fc.in_features resnet50.fc = nn.Sequential( nn.Linear(fc_inputs, 256), nn.ReLU(), nn.Dropout(0.4), nn.Linear(256, 2), nn.LogSoftmax(dim=1) ) # 用GPU进行训练。 resnet50 = resnet50.to('cpu') # 定义损失函数和优化器。 loss_func = nn.CrossEntropyLoss() optimizer = optim.Adam(resnet50.parameters()) # 五、训练 def train_and_valid(model, loss_function, optimizer, epochs=25): device = torch.device("cuda:1" 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(tqdm(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(tqdm(test_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 / test_data_size avg_valid_acc = valid_acc / test_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_SAVE_PATH = "./" MODEL_NAME = 'models/'+dataset+'_model_'+str(epoch+1)+'.pt' torch.save(model, os.path.join(MODEL_SAVE_PATH, MODEL_NAME)) return model, history num_epochs = 30 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, 1) 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()
