import torch import torchvision.models as models # 导入存有各种深度学习模型的模块 from torch import nn # 导入神经网络模块 from torch.utils.data import Dataset, DataLoader # Dataset： 抽象类，一种用于获取数据的方法 DataLoader：数据包管理工具，打包数据 from torchvision import transforms # transforms模块提供了一系列用于图像预处理和数据增强的函数和类 from PIL import Image # 用于处理图片 import numpy as np """ 调用resnet18模型 """ resnet_model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT) for param in resnet_model.parameters(): param.requires_grad = False # 模型所有参数(即权重和偏差)的 requires_grad 属性设置成 False，从而冻结所有模型参数 # 使得在反向传播过程中不会计算他们的梯度，从此减少模型的计算量，提高推理速度 in_features = resnet_model.fc.in_features # 获取resnet18模型全连接层原输入的特征个数 resnet_model.fc = nn.Linear(in_features, 20) # 创建一个全连接层输入特征个数为： in_features 输出特征个数为：数据集中事务的种类数量 params_to_update = [] # 保存需要训练的参数，仅仅包含全连接层的参数 for param in resnet_model.parameters(): if param.requires_grad == True: params_to_update.append(param) """ 图像预处理和数据增强 """ data_transforms = { 'train': transforms.Compose([ transforms.Resize([300, 300]), transforms.RandomRotation(45), # 随机旋转，-45到45度之间随机 transforms.CenterCrop(224), # 中心裁剪 transforms.RandomHorizontalFlip(p=0.5), # 随机水平反转 选择一个概率 transforms.RandomVerticalFlip(p=0.5), # 随机垂直翻转 # transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1), # 亮度、对比度 transforms.RandomGrayscale(p=0.1), # 概率转换成灰度率，3通道就是R G B transforms.ToTensor(), # 转化为神经网络可以识别的 Tensor 类型 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 对图片数据进行归一化，[均值]，[标准差] ]), 'valid': transforms.Compose([ transforms.Resize([224, 224]), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } """ 定义获取每张食物图片和标签的类方法 """ class food_dataset(Dataset): def __init__(self, file_path, transform=None): self.file_path = file_path self.imgs = [] self.labels = [] self.transform = transform with open(self.file_path) as f: samples = [x.strip().split(' ') for x in f.readlines()] for img_path, label in samples: self.imgs.append(img_path) self.labels.append(label) def __len__(self): return len(self.imgs) def __getitem__(self, idx): image = Image.open(self.imgs[idx]) if self.transform: image = self.transform(image) label = self.labels[idx] label = torch.from_numpy(np.array(label, dtype=np.int64)) return image, label """ 获取训练集和测试集数据 """ training_data = food_dataset(file_path='trainda.txt', transform=data_transforms['train']) test_data = food_dataset(file_path='testda.txt', transform=data_transforms['valid']) """ 对数据集进行打包 """ train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True) # 64张图片为一个包，shuffle --> 打乱顺序 test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True) """ 判断当前设备是否支持GPU，其中mps是苹果m系列芯片的GPU """ device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" print(f"Using {device} device") # 把模型传入到 gpu 或 cpu model = resnet_model.to(device) """ 调用交叉熵损失函数 """ loss_fn = nn.CrossEntropyLoss() """" 创建优化器并调整优化器中的学习率--> lr """ optimizer = torch.optim.Adam(params_to_update, lr=0.001) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5) """ 定义训练模型的函数 """ def train(dataloader, model, loss_fn, optimizer): model.train() # 告诉模型，开始训练 # pytorch提供2种方式来切换训练和测试的模式，分别是:model.train()和 model.eval()。 # 一般用法是:在训练开始之前写上model.trian()，在测试时写上model. for X, y in dataloader: X, y = X.to(device), y.to(device) # 把训练数据集和标签传入cpu或gpu pred = model.forward(X) # .forward可以被省略，父类中已经对次功能进行了设置。自动初始化w权值 loss = loss_fn(pred, y) # 通过交叉熵损失函数计算损失值 loss optimizer.zero_grad() # 梯度值清零 loss.backward() # 反向传播计算得到每个参数的梯度值w optimizer.step() # 根据梯度更新网络w参数 """ 定义测试模型的函数 """ best_acc = 0 # 用于更新准确率 def test(dataloader, model, loss_fn): global best_acc size = len(dataloader.dataset) num_batches = len(dataloader) model.eval() # 测试，w就不能再更新 test_loss, correct = 0, 0 with torch.no_grad(): for X, y in dataloader: X, y = X.to(device), y.to(device) pred = model.forward(X) test_loss += loss_fn(pred, y).item() # test_loss是会自动累加每一个批次的损失值 correct += (pred.argmax(1) == y).type(torch.float).sum().item() # correct是会自动累加每一个批次的正确率 test_loss /= num_batches # 平均的损失值 correct /= size # 平均的正确率 print(f"Test result: \n Accuracy: {(100 * correct)}%, Avg loss: {test_loss}") # 找到最好的准确率 if correct > best_acc: best_acc = correct """ 定义模型训练的轮数，并每训练一轮测试一次 """ epochs = 30 for e in range(epochs): print(f"Epoch {e + 1}\n---------------------------") train(train_dataloader, model, loss_fn, optimizer) scheduler.step() # 在每个epoch的训练中，使用scheduler.step()语句进行学习率更新 test(test_dataloader, model, loss_fn) print('最优的训练结果为：', best_acc)