多的不说，直接放码过来：

vit的主要思想就是将图片切割为多个patch块，大小为patch_size,数量为(size/patch_size)^2

对每个patch展平为一维向量，传入transformer的编码器中得到提取特征后的向量，这样就和nlp里面的任务一样了！

import torch

import torch.nn as nn

import torchvision

from torch.utils.data import DataLoader

from torchvision import transforms

from torch import optim

import timeit

from tqdm import tqdm

class PatchEmbedding(nn.Module):

    def __init__(self, in_channels, patch_size, embed_dim, num_patches, dropout):

        super(PatchEmbedding, self).__init__()

        self.patcher = nn.Sequential(

            nn.Conv2d(in_channels=in_channels, out_channels=embed_dim, kernel_size=patch_size, stride=patch_size),

            nn.Flatten(2)

        )

        self.cls_token = nn.Parameter(torch.randn(size=(1, 1, embed_dim)), requires_grad=True)

        self.postion_embedding = nn.Parameter(torch.randn(size=(1, num_patches + 1, embed_dim)), requires_grad=True)

        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x):

        # cls_token = self.cls_token.expand(x.shape[0], -1, -1)

        cls_token = self.cls_token.expand(x.shape[0], 1, -1)  # [batch_size,1,768(embed_dim)]

        x = self.patcher(x).permute(0, 2, 1)  # [batch_size,patches,embed_dim]

        x = torch.cat([cls_token, x], dim=1)  # 拼接分类编码

        x = x + self.postion_embedding  # size=(1, num_patches + 1, embed_dim)

        x = self.dropout(x)

        return x

class Vit(nn.Module):

    def __init__(self, in_channels, patch_size, embed_dim, num_patches, dropout,

                 num_head, activation, num_encoders, num_class):

        super(Vit, self).__init__()

        self.patch_embedding = PatchEmbedding(in_channels, patch_size, embed_dim, num_patches, dropout)

        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_head, dropout=dropout,

                                                   activation=activation, batch_first=True, norm_first=True)

        # 使用多个TransformerEncoderLayer实例化TransformerEncoder

        self.encoder_layers = nn.TransformerEncoder(encoder_layer, num_layers=num_encoders)

        self.MLP = nn.Sequential(

            nn.LayerNorm(normalized_shape=embed_dim),

            nn.Linear(in_features=embed_dim, out_features=num_class)

        )

    def forward(self, x):

        x = self.patch_embedding(x)

        x = self.encoder_layers(x)

        x = self.MLP(x[:, 0, :])  # 取cls_token

        return x

in_channels = 3

img_size = 32

patch_size = 8

embed_dim = patch_size ** 2 * in_channels

num_patches = (img_size // patch_size) ** 2

dropout = 0.01

batch_size = 64

device = "cuda" if torch.cuda.is_available() else "cpu"

epochs = 50

num_head = 8

activation = "gelu"

num_encoders = 10

num_classes = 10

learning_rate = 1e-4

weight_dacay = 1e-4

betas = (0.9, 0.999)

train_transform = torchvision.transforms.Compose([

    # transforms.ToPILImage(),

    transforms.RandomRotation(15),

    transforms.ToTensor(),

    transforms.Normalize(

        mean=0.5, std=0.5

    )

])

test_transform = torchvision.transforms.Compose([

    # transforms.ToPILImage(),

    # transforms.RandomRotation(15),

    transforms.ToTensor(),

    transforms.Normalize(

        mean=0.5, std=0.5

    )

])

train_dataset = torchvision.datasets.CIFAR10(root="../../datas", train=True, transform=train_transform)

train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = torchvision.datasets.CIFAR10(root="../../datas", train=False, transform=test_transform)

test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

model = Vit(in_channels, patch_size, embed_dim, num_patches, dropout,

            num_head, activation, num_encoders, num_classes).to(device)

criterion = nn.CrossEntropyLoss()

optimzer = optim.Adam(model.parameters(), lr=learning_rate, betas=betas, weight_decay=weight_dacay)

start = timeit.default_timer()

best_acc = 0

print(f"training on : {device}")

for epoch in range(epochs):

    model.train()

    train_labels = []

    train_preds = []

    train_running_loss = 0

    n = 0

    train_n_sum = 0

    train_n_correct = 0

    for idx, (X, y) in enumerate(tqdm(train_dataloader, position=0, leave=True)):

        X = X.to(device)

        y = y.to(device)

        y_pred = model(X)

        y_pred_label = torch.argmax(y_pred, dim=1)

        # print(y.shape, y_pred.shape)

        train_labels.extend(y.cpu().detach())

        train_preds.extend(y_pred.cpu().detach())

        loss = criterion(y_pred, y)

        optimzer.zero_grad()

        loss.backward()

        optimzer.step()

        train_running_loss += loss.item()

        n += 1

        train_n_sum += X.size(0)

        train_n_correct += (y == y_pred_label).sum().item()

    train_loss = train_running_loss / (n + 1)

    train_acc = train_n_correct / train_n_sum

    model.eval()

    val_labels = []

    val_preds = []

    val_running_loss = 0

    test_n_sum = 0

    test_n_correct = 0

    with torch.no_grad():

        n = 0

        for idx, (X, y) in enumerate(tqdm(test_dataloader, position=0, leave=True)):

            X = X.to(device)

            y = y.to(device)

            y_pred = model(X)

            y_pred_label = torch.argmax(y_pred, dim=1)

            # print(y_pred.shape, y.shape)

            val_labels.extend(y.cpu().detach())

            val_preds.extend(y_pred.cpu().detach())

            loss = criterion(y_pred, y)

            val_running_loss += loss.item()

            n += 1

            test_n_sum += X.size(0)

            test_n_correct += (y == y_pred_label).sum().item()

        test_loss = val_running_loss / (n + 1)

        test_acc = test_n_correct / test_n_sum

        if test_acc > best_acc:

            best_acc = test_acc

            map = {

                'state': model.state_dict(),

                "acc": test_acc,

                "loss": test_loss

            }

            print("save model : ", map['acc'])

            torch.save(map, "./checkpoints/vit_model.pth")

    print("-" * 30)

    print(f"train loss epoch : {epoch + 1} : {train_loss:.4f}")

    print(f"test loss epoch : {epoch + 1} : {test_loss:.4f}")

    print(

        f"train acc epoch : {epoch + 1} : {train_acc:.4f}"

    )

    print(

        f"test acc epoch : {epoch + 1} : {test_acc:.4f}"

    )

    print("-" * 30)

stop = timeit.default_timer()

print(f"training time : {stop - start:.2f}")

# patcher = PatchEmbedding(in_channels=in_channels, patch_size=patch_size,

#                          embed_dim=embed_dim, num_patches=num_patches, dropout=dropout)

# for idx, (x, y) in enumerate(dataloader):

#     print("before : ", idx, x.shape, y.shape)

#     x = patcher(x)

#     print("after : ", x.shape)

#     break

附带几个训练效果图：