transformers做cifar10图像分类任务效果好保存训练好的模型且有loss变化曲线以及准确率变化曲线的代码

以下是使用transformers库进行CIFAR-10图像分类任务的代码，包括保存训练好的模型以及绘制loss变化曲线和准确率变化曲线：

import torch
import torchvision
import torch.optim as optim
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import DataLoader
from sklearn.metrics import accuracy_score
from transformers import ViTModel, ViTConfig

# 定义ViT模型
class ViT(torch.nn.Module):
    def __init__(self):
        super(ViT, self).__init__()
        self.vit_model = ViTModel(ViTConfig(image_size=32, patch_size=4, num_classes=10))

    def forward(self, img):
        img = self.vit_model.pixel_embed(img)
        img = self.vit_model.transformer(img)
        img = self.vit_model.ln_post(img)
        img = img[:, 0, :]
        return img

# 定义训练函数
def train(model, optimizer, train_loader, device):
    model.train()
    train_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        train_loss += loss.item()
        loss.backward()
        optimizer.step()
    return train_loss / len(train_loader)

# 定义测试函数
def test(model, test_loader, device):
    model.eval()
    test_loss = 0
    predictions, true_labels = [], []
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = F.cross_entropy(output, target)
            test_loss += loss.item()
            pred = output.argmax(dim=1, keepdim=True)
            predictions += pred.cpu().numpy().tolist()
            true_labels += target.cpu().numpy().tolist()
    acc = accuracy_score(true_labels, predictions)
    return test_loss / len(test_loader), acc

# 定义数据增强
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
])

# 加载CIFAR-10数据集
train_data = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
test_data = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)

batch_size = 128
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=2)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=2)

# 定义模型和优化器
model = ViT().cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练和测试模型
epochs = 50
train_losses, test_losses, accs = [], [], []
for epoch in range(1, epochs + 1):
    train_loss = train(model, optimizer, train_loader, 'cuda')
    test_loss, acc = test(model, test_loader, 'cuda')
    train_losses.append(train_loss)
    test_losses.append(test_loss)
    accs.append(acc)
    print('Epoch: {}, Train Loss: {:.6f}, Test Loss: {:.6f}, Acc: {:.6f}'.format(epoch, train_loss, test_loss, acc))

# 保存模型
torch.save(model.state_dict(), 'vit_cifar10.pth')

# 绘制loss变化曲线和准确率变化曲线
import matplotlib.pyplot as plt

plt.plot(train_losses)
plt.plot(test_losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend(['Train Loss', 'Test Loss'])
plt.show()

plt.plot(accs)
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.show()

其中，模型使用的是ViT模型，优化器使用的是Adam，epochs为50，训练过程中保存了train loss、test loss和accuracy，训练完成后将模型保存为'vit_cifar10.pth'，并绘制了loss变化曲线和准确率变化曲线。