# 1.导入工具包 import torch import torchvision.models as models # 2.实例化模型 # model = models.resnet18() model = models.resnet50() # print(model) # 3.查看模型结构 # print(model.layer4) # --------------------------------鲜花识别案例---------------------------------------- # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 1.读取数据 from torchvision import transforms from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder # 1.1 指定批次大小 batch_size = 256 # 1.2 指定数据集路径 flower_train_path = './dataset/flower_datas/train/' flower_test_path = './dataset/flower_datas/val/' # 1.3 先将数据转换为tensor类型，并调整数据的大小为224x224 dataset_transform = transforms.Compose( [transforms.ToTensor(), transforms.Resize((224,224))]) # 1.4 获取训练集数据和测试集数据 flower_train = ImageFolder(flower_train_path, transform=dataset_transform) flower_test = ImageFolder(flower_train_path, transform=dataset_transform) # 1.5 获取数据的迭代 train_loader = DataLoader(dataset=flower_train, batch_size=batch_size, shuffle=True) test_loader = DataLoader(dataset=flower_test, batch_size=batch_size, shuffle=False) # import matplotlib.pyplot as plt # 1.6 遍历每个迭代的数据，将其结果展示出来 # for b, (imgs, targets) in enumerate(train_loader): # # 获取第一个batch的图像 # if b == 10: # # 将其进行展示 # fig, axes = plt.subplots(1, 2) # # 遍历batch中的每个图像 # for i in range(batch_size): # # 图像显示出来 # axes[i].imshow(imgs[i].permute(1, 2, 0)) # # 设置图像标题 # axes[i].set_title(targets[i].item()) # plt.show() # elif b > 10: # break # 模型实例化:输入数据3通道，进行5类的分类处理 # 2.模型实例化和参数设置 # 2.1 模型实例化:输入数据3通道，进行5类的分类处理 model = models.resnet18(num_classes=5).to(device) # 2.2 模型训练的参数设置 # 2.2.1 学习率 learning_rate = 1e-3 # 2.2.2 训练轮数 num_epochs = 10 # 2.2.3 优化算法Adam = RMSProp + Momentum optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # 2.2.4 交叉熵损失函数 loss_fn = torch.nn.CrossEntropyLoss() # 3.定义评估函数 def evaluate_accuracy(data_iter, model): total = 0 correct = 0 # 3.1 不进行梯度计算 with torch.no_grad(): # 3.2 模型是验证模式 model.eval() # 3.3 获取每一个batch的数据，进行预测 for images, labels in data_iter: images,labels = images.to(device),labels.to(device) outputs = model(images).to(device) # 3.4 获取预测结果 _, predicts = torch.max(outputs.data, dim=1) # 3.5 预测的次数 total += labels.size(0) # 3.6 预测正确的个数 correct += (predicts == labels).cpu().sum() # 3.7 获取准确率 return correct / total # 4.定义模型训练过程:指定数据集，优化器，损失函数和轮次 def train(data_loader=train_loader, optimizer=optimizer, loss_fn=loss_fn, epochs=num_epochs): # 4.1 遍历每一个轮次进行训练 for epoch in range(epochs): print('current epoch = {}'.format(epoch)) # 4.2 定义初试参数 # 每一个轮次的损失，预测个数和预测正确个数的初始化 train_accuracy_total = 0 train_correct = 0 # 损失值的和 train_loss_sum = 0 # 迭代次数 iter = 0 # 4.3 遍历每一个batch的数据进行训练 for i, (images, labels) in enumerate(data_loader): images,labels = images.to(device),labels.to(device) # 4.4 模型定义为训练模式 model.train() # 4.5 对数据进行预测 outputs = model(images).to(device) # 4.6 计算模型的损失 loss = loss_fn(outputs, labels) # 4.7 在做反向传播前先清除网络状态 optimizer.zero_grad() # 4.8 损失值进行反向传播 loss.backward() # 4.9 参数迭代更新 optimizer.step() # 4.10 求损失的和 train_loss_sum += loss.item() # 4.11 输出模型预测结果 _, predicts = torch.max(outputs.data, dim=1) # 4.12 获取训练集预测正确的个数 train_accuracy_total += labels.size(0) train_correct += (predicts == labels).cpu().sum().item() iter+=1 # 4.13 测试集预测的准确率 test_acc = evaluate_accuracy(test_loader, model) print( 'epoch:{0}, loss:{1:.4f}, train accuracy:{2:.3f}, test accuracy:{3:.3f}' .format(epoch, train_loss_sum / (iter+0.01), train_correct / train_accuracy_total, test_acc)) print('------------finish training-------------') if __name__ == '__main__': train()