本文介绍了使用PyTorch进行图像分类的过程，包括数据集的准备，如从下载的flower_photos.zip中划分训练集和验证集；模型训练，以ResNet101为例，预训练模型并进行参数调整；以及模型的训练、验证和推理，包括损失函数、优化器选择和模型精度评估。
提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档
前言
提示：这里可以添加本文要记录的大概内容：
例如：随着人工智能的不断发展，机器学习这门技术也越来越重要，很多人都开启了学习机器学习，本文就介绍了机器学习的基础内容。
提示：以下是本篇文章正文内容，下面案例可供参考
一、数据集准备
需要把数据拆分为训练集和验证集
二、模型训练
1.训练
代码如下（示例）：
import time
import torch
import os
from torchvision import datasets,transforms
import torch.optim as optim
from torch import nn
from torch.utils.data import DataLoader
from torchvision.models import alexnet # 最简单的模型
from torchvision.models import vgg11, vgg13, vgg16, vgg19 # VGG系列
from torchvision.models import resnet18, resnet34,resnet50, resnet101, resnet152 # ResNet系列
from torchvision.models import inception_v3
from torchvision.models import resnet18, resnet34,resnet50, resnet101, resnet152
def pre_model(pretrained):
model = resnet101(pretrained=False, num_classes=5)
if pretrained:
model = resnet101(pretrained=True)
for param in model.parameters():
param.requires_grad = False#上面所有的参数保持不动
in_f = model.fc.in_features#保证最后这个参数准确的传入
model.fc = nn.Linear(in_f, 5)#替换掉最后的fc的层，分成4种类型
for param in model.fc.parameters():
param.requires_grad = True
return model
data_transforms = {
'train': transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'val': transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
}
#从指定路径获取数据集
train_path = "./dataset/flower_data/train"
val_path = "./dataset/flower_data/val"
train_data = datasets.ImageFolder(train_path, data_transforms["train"])
val_data = datasets.ImageFolder(val_path, data_transforms["val"])
print("train_data:", train_data)
print("val_data:", val_data)
# 获取所有类别名称
classes = train_data.classes
labels = ','.join(classes)
print("train_data.class: ", labels)
# 把分类标签写入到label.txt文件中
with open(os.path.join("./",'label.txt'), 'w', encoding='utf-8') as file:
file.write(labels)
# 创建三个加载器，分别为训练，验证，将训练集的batch大小设为64，即每次加载器向网络输送64张图片，随机打乱
trainloader = torch.utils.data.DataLoader(train_data,batch_size = 64,shuffle = True)
validloader = torch.utils.data.DataLoader(val_data,batch_size = 20)
# 如果显卡可用，则用显卡进行训练
device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)
model = pre_model(pretrained=True)
model.to(device)
print(model)
# 定义损失函数，计算相差多少，交叉熵，
loss_fn = nn.CrossEntropyLoss()
# 定义优化器，用来训练时候优化模型参数，随机梯度下降法
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) # 初始学习率
# 学习率衰减，每7 epochs衰减0.1
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)
# 一共训练1次
epochs = 100
size = len(validloader.dataset)
trainloader_size = len(trainloader.dataset)
print("validloader size: ", size)
best_acc = 0
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
time_start = time.time()
test_loss, correct = 0, 0
# 训练模式
model.train()
# 从数据加载器中读取batch（一次读取多少张，即批次数），X(图片数据)，y（图片真实标签）。
for batch, (image, label) in enumerate(trainloader):
# 将数据存到显卡
image, label = image.cuda(), label.cuda()
# 得到预测的结果pred
outputs = model(image)
_, preds = torch.max(outputs, 1)#取出我们的预测最大值的索引 用来计算后面的准确率
# 计算预测的误差
#print(pred,label)
loss = loss_fn(outputs,label)
# 反向传播，更新模型参数
optimizer.zero_grad()
loss.backward()
optimizer.step()
test_loss += loss.item() * image.size(0) # 计算损失函数的平均值
correct += torch.sum(preds == label.data)
# 每训练10次，输出一次当前信息
if batch % 10 == 0:
current = batch * len(image)
# 统计预测正确的个数
print(f"train, total loss: {test_loss:>8f} [{current:>5d}/{trainloader_size:>5d}] \n")
time_end = time.time()
print(f"train time: {(time_end-time_start)}")
loss = test_loss / trainloader_size
acc = correct.double() / trainloader_size
print(f"train: Accuracy: {(100*acc):>0.1f}%, Avg loss: {loss:>8f} \n")
scheduler.step()
# 将模型转为验证模式
model.eval()
# 初始化test_loss 和 correct， 用来统计每次的误差
test_loss, correct = 0, 0
total = 0
# 测试时模型参数不用更新，所以no_gard()
# 禁用梯度，非训练， 推理期用到
with torch.no_grad():
# 加载数据加载器，得到里面的X（图片数据）和y(真实标签）
for image, label in validloader:
# 将数据转到GPU
image, label = image.cuda(), label.cuda()
# 将图片传入到模型当中就，得到预测的值pred
outputs = model(image)
_, preds = torch.max(outputs, 1)#取出我们的预测最大值的索引 用来计算后面的准确率
# 计算预测值pred和真实值y的差距
loss = loss_fn(outputs,label)
test_loss += loss.item() * image.size(0) # 计算损失函数的平均值
correct += torch.sum(preds == label.data)
loss = test_loss / size
acc = correct / size
print(f"val: correct = {correct}, Test Error: \n Accuracy: {(100*acc):>0.1f}%, Avg loss: {loss:>8f} \n")
if acc>best_acc:
best_acc=acc
print(f'best_acc ({best_acc:.6f}--->{acc:.6f}) \t Saving The Model')
#torch.save(model.state_dict(), os.path.join("./", 'best_model.pth'))
torch.save(model.fc.state_dict(), os.path.join("./", 'best_head_model.pth'))
print("Done!")
2.推理
代码如下（示例）：
from torchvision.models import resnet18, resnet34,resnet50, resnet101, resnet152 from PIL import Image import os import torch from torch import nn from torchvision import transforms import numpy as np transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) device = "cuda" if torch.cuda.is_available() else "cpu" path = "./dataset/flower_data/train/roses/16525204061_9b47be3726_m.jpg" image = Image.open(path) image.convert("RGB") transformed_image = transform(image) print("transformed_image: ", transformed_image.shape) input_tensor = transformed_image.unsqueeze(0).to(device) print("input_tensor: ", input_tensor.shape) with open(os.path.join("./",'label.txt'), 'r') as file: content = file.read() classes = content.split(',') print("classes: ", classes) model = resnet101(pretrained=True).to(device) in_f = model.fc.in_features#保证最后这个参数准确的传入 model.fc = nn.Linear(in_f, 5)#替换掉最后的fc的层，分成4种类型 # 加载训练的分类器权重 state_dict = torch.load("best_head_model.pth") model.fc.load_state_dict(state_dict) model.to(device) model.eval() with torch.no_grad(): result = model(input_tensor) print("result:", result) predicted = result.detach().cpu().numpy() print("predicted:", predicted) result = np.argmax(predicted) print("result:", result) print(classes[result])
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
