深度学习（VIT）
将Transformer引入图像领域之作，学习一下。
网络结构：
VIT结构有几个关键的地方：
1. 图像分块：输入图像被划分为固定大小的非重叠小块（patches），每个小块被展平并线性嵌入到一个固定维度的向量中。这里是将32x32的图像划分成4x4的小块，总共会有16个小块，每个小块有64维向量。
2. 位置编码：由于Transformer不具备位置敏感性，需要添加位置编码来提供位置信息。每个图像块向量都会加上一个对应的可学习的位置编码，以保留图像空间信息。
3. Transformer编码：嵌入向量连同位置编码一起被输入到Transformer编码器中，编码器由多个相同自注意力层堆叠而成。
4. MLP分类。
测试代码如下：
import torch import torch.nn as nn import torch.optim as optim from torchvision import transforms from torchvision.datasets import CIFAR10 import torchvision.models class EmbedLayer(nn.Module): def __init__(self,channels, embed_dim,img_size,patch_size): super().__init__() self.embed_dim = embed_dim self.conv1 = nn.Conv2d(channels, embed_dim, patch_size, patch_size) self.pos_embedding = nn.Parameter(torch.zeros(1, (img_size // patch_size) ** 2,embed_dim), requires_grad=True) # Positional Embedding def forward(self, x): x = self.conv1(x) x = x.reshape([x.shape[0], self.embed_dim, -1]) x = x.transpose(1, 2) x = x + self.pos_embedding return x class SelfAttention(nn.Module): def __init__(self,embed_dim, heads): super().__init__() self.heads = heads self.embed_dim = embed_dim self.head_embed_dim = self.embed_dim // heads self.queries = nn.Linear(self.embed_dim, self.head_embed_dim * heads, bias=True) self.keys = nn.Linear(self.embed_dim, self.head_embed_dim * heads, bias=True) self.values = nn.Linear(self.embed_dim, self.head_embed_dim * heads, bias=True) def forward(self, x): m, s, e = x.shape q = self.queries(x).reshape(m, s, self.heads, self.head_embed_dim).transpose(1, 2) k = self.keys(x).reshape(m, s, self.heads, self.head_embed_dim).transpose(1, 2) v = self.values(x).reshape(m, s, self.heads, self.head_embed_dim).transpose(1, 2) q = q.reshape([-1, s, self.head_embed_dim]) k = k.reshape([-1, s, self.head_embed_dim]) v = v.reshape([-1, s, self.head_embed_dim]) k = k.transpose(1, 2) x_attention = q.bmm(k) x_attention = torch.softmax(x_attention, dim=-1) x = x_attention.bmm(v) x = x.reshape([-1, self.heads, s, self.head_embed_dim]) x = x.transpose(1, 2) x = x.reshape(m, s, e) return x class Encoder(nn.Module): def __init__(self, embed_dim,heads): super().__init__() self.attention = SelfAttention(embed_dim,heads) self.fc1 = nn.Linear(embed_dim, embed_dim * 2) self.activation = nn.GELU() self.fc2 = nn.Linear(embed_dim * 2,embed_dim) self.norm1 = nn.LayerNorm(embed_dim) self.norm2 = nn.LayerNorm(embed_dim) def forward(self, x): x = x + self.attention(self.norm1(x)) x = x + self.fc2(self.activation(self.fc1(self.norm2(x)))) return x class Classifier(nn.Module): def __init__(self, embed_dim,num_patches,classes): super().__init__() self.fc1 = nn.Linear(embed_dim*num_patches, embed_dim) self.activation = nn.Tanh() self.fc2 = nn.Linear(embed_dim, classes) def forward(self, x): x = x.view(x.shape[0],-1) x = self.fc1(x) x = self.activation(x) x = self.fc2(x) return x class VisionTransformer(nn.Module): def __init__(self,channels, embed_dim,n_layers,heads,img_size,patch_size,classes): super().__init__() self.embedding = EmbedLayer(channels,embed_dim,img_size,patch_size) self.encoder = nn.Sequential(*[Encoder(embed_dim,heads) for _ in range(n_layers)], nn.LayerNorm(embed_dim)) self.norm = nn.LayerNorm(embed_dim) self.classifier = Classifier(embed_dim,(img_size//patch_size)**2,classes) def forward(self, x): x = self.embedding(x) x = self.encoder(x) x = self.norm(x) x = self.classifier(x) return x if __name__ == '__main__': device = torch.device("cuda") trainTransforms = transforms.Compose([ transforms.ToTensor() , transforms.RandomHorizontalFlip(p=0.5) , transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) testTransforms = transforms.Compose([ transforms.ToTensor() , transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) trainset = CIFAR10(root='./data', train=True, download=True, transform=trainTransforms) trainLoader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True) testset = CIFAR10(root='./data', train=False,download=False, transform=testTransforms) testLoader = torch.utils.data.DataLoader(testset, batch_size=128,shuffle=False) model = VisionTransformer(channels=3, embed_dim=128,n_layers=6,heads=8,img_size=32,patch_size=8,classes=10) # model = torchvision.models.resnet18(pretrained=True) # model.conv1 = nn.Conv2d(3, 64, 3, stride=1, padding=1, bias=False) # model.maxpool = nn.MaxPool2d(1, 1, 0) # model.fc = nn.Linear(model.fc.in_features, 10) criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=5e-4, weight_decay=1e-3) cos_decay = optim.lr_scheduler.CosineAnnealingLR(optimizer, 100, verbose=True) model.to(device) for epoch in range(50): print("epoch :",epoch) model.train() correct = 0 total = 0 for images, labels in trainLoader: images = images.to(device) labels = labels.to(device) outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(loss.item(),f" train Accuracy: {(100 * correct / total):.2f}%") cos_decay.step() model.eval() with torch.no_grad(): correct = 0 total = 0 for images,labels in testLoader: images = images.to(device) labels = labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f"test Accuracy: {(100 * correct / total):.2f}%") # 保存模型 torch.save(model.state_dict(), 'vit.pth')
