transformers做cifar10图像分类有loss曲线和准确率曲线并保存模型pytorch
时间: 2024-02-11 12:08:14 浏览: 135
利用pytorch对CIFAR数据进行图像分类(包含全套代码和10+个模型的实现)
5星 · 资源好评率100%
好的,以下是使用 Transformers 模型进行 CIFAR-10 图像分类并保存模型的 PyTorch 代码,包括训练过程中的损失和准确率曲线:
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from transformers import ViTModel, ViTConfig
# CIFAR-10 数据集
transform_train = transforms.Compose(
[
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
)
transform_test = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
)
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=64,
shuffle=False, num_workers=2)
# 模型定义
config = ViTConfig.from_pretrained('google/vit-base-patch16-224')
model = ViTModel(config)
model.classifier = nn.Linear(config.hidden_size, 10)
# 训练参数
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 训练过程
num_epochs = 10
train_loss = []
test_loss = []
train_acc = []
test_acc = []
for epoch in range(num_epochs):
# 训练
model.train()
running_loss = 0.0
correct = 0
total = 0
for i, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs.logits, targets)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.logits.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
epoch_loss = running_loss / len(trainloader)
epoch_acc = correct / total
train_loss.append(epoch_loss)
train_acc.append(epoch_acc)
# 测试
model.eval()
running_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
for i, (inputs, targets) in enumerate(testloader):
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
loss = criterion(outputs.logits, targets)
running_loss += loss.item()
_, predicted = outputs.logits.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
epoch_loss = running_loss / len(testloader)
epoch_acc = correct / total
test_loss.append(epoch_loss)
test_acc.append(epoch_acc)
print(f"Epoch: {epoch+1}/{num_epochs} | Train Loss: {train_loss[-1]:.4f} | Train Acc: {train_acc[-1]*100:.2f}% | Test Loss: {test_loss[-1]:.4f} | Test Acc: {test_acc[-1]*100:.2f}%")
# 调整学习率
scheduler.step()
# 保存模型
torch.save(model.state_dict(), "vit_cifar10.pth")
```
在上述代码中,我们使用 PyTorch 自带的 CIFAR-10 数据集,并对数据进行了预处理。然后,我们定义了一个 ViT 模型,并将其分类器层替换为一个具有 10 个输出的线性层。接下来,我们定义了训练参数,包括损失函数、优化器、学习率调度器和设备。在训练循环中,我们对模型进行了训练和测试,并记录了训练和测试的损失和准确率。最后,我们将训练好的模型保存到文件中。
你可以根据需要调整代码中的超参数和模型定义来适应你的应用场景。
阅读全文