如何用vits训练模型

ViT（Vision Transformer）是一种基于Transformer的视觉模型，它在图像分类、目标检测和图像分割等任务上取得了很好的性能。下面是一个简单的示例，展示如何使用PyTorch和Imagenet数据集训练一个ViT模型：

1.首先，需要安装 PyTorch 和 torchvision 库。可以使用以下命令安装：

pip install torch torchvision

2.下载并解压Imagnet数据集，可以从官网下载：http://www.image-net.org/

3.接下来，需要准备数据集。可以使用 torchvision.datasets.ImageFolder 类来加载数据集，并使用 torchvision.transforms 模块来对图像进行预处理和数据增强。

import torchvision.datasets as datasets
import torchvision.transforms as transforms

train_dataset = datasets.ImageFolder(
    'path/to/imagenet/train',
    transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.RandomCrop((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ])
)

val_dataset = datasets.ImageFolder(
    'path/to/imagenet/val',
    transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.CenterCrop((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ])
)

4.定义模型

import torch.nn as nn
from timm.models.vision_transformer import vit_base_patch16_224

model = vit_base_patch16_224(pretrained=False)
num_classes = 1000
model.head = nn.Linear(model.head.in_features, num_classes)

5.定义优化器和损失函数

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

6.开始训练模型

from torch.utils.data import DataLoader

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

num_epochs = 10
for epoch in range(num_epochs):
    train_loss = 0.0
    train_acc = 0.0
    val_loss = 0.0
    val_acc = 0.0

    model.train()
    for inputs, labels in train_loader:
        inputs = inputs.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item() * inputs.size(0)
        _, preds = torch.max(outputs, 1)
        train_acc += torch.sum(preds == labels.data)

    train_loss = train_loss / len(train_dataset)
    train_acc = train_acc / len(train_dataset)

    model.eval()
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            loss = criterion(outputs, labels)

            val_loss += loss.item() * inputs.size(0)
            _, preds = torch.max(outputs, 1)
            val_acc += torch.sum(preds == labels.data)

    val_loss = val_loss / len(val_dataset)
    val_acc = val_acc / len(val_dataset)

    print('Epoch: {}/{}'.format(epoch+1, num_epochs))
    print('Training Loss: {:.4f}, Training Acc: {:.4f}'.format(train_loss, train_acc))
    print('Validation Loss: {:.4f}, Validation Acc: {:.4f}'.format(val_loss, val_acc))

7.保存模型

torch.save(model.state_dict(), 'path/to/save/model.pt')

这是一个简单的示例，可以根据自己的需求进行修改和扩展。