resnet18图像分类

网上搜的代码，加了个损失曲线可视化，需要改下，保存最优模型，懒得改，没力气，代码就是拿着搜到的代码改来改去，感觉结果还不错，数据集再多点应该更好。
import os
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from torchvision import models, datasets, transforms
import torch.utils.data as tud
import numpy as np
from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
from PIL import Image
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings(“ignore”)

device = torch.device(“cuda:0” if torch.cuda.is_available() else ‘cpu’)
n_classes = 3 # 几种分类的
preteain = False # 是否下载使用训练参数 有网true 没网false
epoches = 100 # 训练的轮次
traindataset = datasets.ImageFolder(root=‘F:/cloud/resnet18/dataset/train’, transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

]))

testdataset = datasets.ImageFolder(root=‘F:/cloud/resnet18/dataset/test’, transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

]))

classes = testdataset.classes
print(classes)

model = models.resnet18(pretrained=preteain)
if preteain == True:
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Linear(in_features=512, out_features=n_classes, bias=True)
model = model.to(device)

def train_model(model, train_loader, loss_fn, optimizer, epoch):
model.train()
total_loss = 0.
total_corrects = 0.
total = 0.
for idx, (inputs, labels) in enumerate(train_loader):
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = loss_fn(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
preds = outputs.argmax(dim=1)
total_corrects += torch.sum(preds.eq(labels))
total_loss += loss.item() * inputs.size(0)
total += labels.size(0)

total_loss = total_loss / total
acc = 100 * total_corrects / total
print("轮次:%4d|训练集损失:%.5f|训练集准确率:%6.2f%%" % (epoch + 1, total_loss, acc))
# 保存模型参数
torch.save(model.state_dict(), 'F:/cloud/resnet18/logs/log.pth')
return total_loss, acc

def test_model(model, test_loader, loss_fn, optimizer, epoch):
model.train()
total_loss = 0.
total_corrects = 0.
total = 0.
with torch.no_grad():
for idx, (inputs, labels) in enumerate(test_loader):
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = loss_fn(outputs, labels)
preds = outputs.argmax(dim=1)
total += labels.size(0)
total_loss += loss.item() * inputs.size(0)
total_corrects += torch.sum(preds.eq(labels))

    loss = total_loss / total
    accuracy = 100 * total_corrects / total
    print("轮次:%4d|测试集损失:%.5f|测试集准确率:%6.2f%%" % (epoch + 1, loss, accuracy))
    return loss, accuracy

loss_fn = nn.CrossEntropyLoss().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.0001)
train_loader = DataLoader(traindataset, batch_size=2, shuffle=True)
test_loader = DataLoader(testdataset, batch_size=2, shuffle=True)

def visualize_losses(train_losses, test_losses, save_path):
plt.plot(train_losses, label=‘Train Loss’)
plt.plot(test_losses, label=‘Test Loss’)
plt.xlabel(‘Epoch’)
plt.ylabel(‘Loss’)
plt.title(‘Train and Test Loss over Epochs’)
plt.legend()
plt.savefig(save_path) # 保存损失曲线为PNG文件
plt.close()

Example usage

train_losses = []
test_losses = []
for epoch in range(0, epoches):
loss1, acc1 = train_model(model, train_loader, loss_fn, optimizer, epoch)
loss2, acc2 = test_model(model, test_loader, loss_fn, optimizer, epoch)
train_losses.append(loss1)
test_losses.append(loss2)

# Visualize train and test losses
save_path = 'F:/cloud/resnet18/logs/loss_curve.png'  # 设置损失曲线保存路径
visualize_losses(train_losses, test_losses, save_path)