写一个用pytorch,VGG16分析JAFFE数据的代码,划分训练集,验证集和测试集,给出准确率图,损失图和混淆矩阵
时间: 2023-06-10 15:03:27 浏览: 193
使用PyTorch将文件夹下的图片分为训练集和验证集实例
3星 · 编辑精心推荐
好的,我理解您的要求了。对于您的问题,我建议您参考以下的代码:
```
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
# 定义超参数
num_epochs = 10
num_classes = 7
batch_size = 16
learning_rate = 0.001
# 加载数据集
transform = transforms.Compose(
[transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])])
train_dataset = torchvision.datasets.ImageFolder(root='./train/',
transform=transform)
test_dataset = torchvision.datasets.ImageFolder(root='./test/',
transform=transform)
# 划分数据集
train_size = int(0.8 * len(train_dataset))
val_size = len(train_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size, shuffle=True)
val_loader = torch.utils.data.DataLoader(dataset=val_dataset,
batch_size=batch_size, shuffle=False)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size, shuffle=False)
# 定义 VGG16 模型
vgg = torchvision.models.vgg16(pretrained=True)
for param in vgg.parameters():
param.requires_grad = False
vgg.classifier[6] = nn.Sequential(
nn.Linear(4096, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, num_classes),
nn.LogSoftmax(dim=1))
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(vgg.classifier[6].parameters(), lr=learning_rate)
# 训练模型并保存最优模型
total_step = len(train_loader)
val_loss_min = np.Inf
for epoch in range(num_epochs):
running_loss = 0.0
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
# 前向传播
outputs = vgg(images)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
if (i + 1) % 10 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))
# 每轮训练后在验证集上计算准确率和损失
val_loss = 0
accuracy = 0
with torch.no_grad():
for images, labels in val_loader:
images = images.to(device)
labels = labels.to(device)
outputs = vgg(images)
val_loss += criterion(outputs, labels).item()
_, predicted = torch.max(outputs.data, 1)
accuracy += (predicted == labels).sum().item()
val_loss /= len(val_loader)
accuracy /= len(val_loader)
print('Epoch [{}/{}], Loss: {:.4f}, Val Loss: {:.4f}, Val Acc: {:.2f}%'.format(epoch + 1, num_epochs, running_loss / total_step, val_loss, accuracy * 100))
# 保存最优模型
if val_loss <= val_loss_min:
print("Validation loss decreased ({:.6f} --> {:.6f}). Saving model ...".format(val_loss_min, val_loss))
torch.save(vgg.state_dict(), 'best_model.pt')
val_loss_min = val_loss
# 加载最优模型并在测试集上计算准确率和混淆矩阵
vgg.load_state_dict(torch.load('best_model.pt'))
with torch.no_grad():
correct = 0
total = 0
y_true = []
y_pred = []
for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
outputs = vgg(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
y_true += labels.tolist()
y_pred += predicted.tolist()
print('Test Accuracy of the model on the {} test images: {:.2f}%'.format(total, (correct / total) * 100))
from sklearn.metrics import confusion_matrix
print(confusion_matrix(y_true, y_pred))
```
在这个代码中,我们使用 VGG16 来对 JAFFE 数据集进行分类,并且将数据集划分成了训练集、验证集和测试集,然后用 pytorch 搭建了 VGG16 模型,并对其进行训练,最后计算了在测试集上的准确率和混淆矩阵。您可以根据自己的需要来修改这个代码,并添加您所需要的准确率图和损失图等。
阅读全文