【论文复现】Vision Transformer（ViT） - ToB企服应用市场:ToB评测及商务社交产业平台

import torch
from vit_pytorch import ViT # 抽象出了一个VIT类
v = ViT(
image_size = 256,# 图片像素大小
patch_size = 32, # patch的大小
num_classes = 1000, # 分类数量
dim = 1024,# 维度
depth = 6, # transformer的block数量
heads = 16, # 线性变换后输出张量的最后维数,多头注意力层中的头数
mlp_dim = 2048, # MLP前馈层维度
dropout = 0.1, # 每个训练步骤中被关闭神经元的比例，可以调成0
emb_dropout = 0.1 # 嵌入丢失率
)
img = torch.randn(1, 3, 256, 256)
preds = v(img) # (1, 1000)

复制代码

import torch
from torch import nn
from einops import rearrange, repeat
from einops.layers.torch import Rearrange
# helpers
def pair(t):
return t if isinstance(t, tuple) else (t, t)
# classes
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout = 0.):
super().__init__()
self.net = nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)
class Attention(nn.Module):
def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
super().__init__()
inner_dim = dim_head * heads
project_out = not (heads == 1 and dim_head == dim)
self.heads = heads
self.scale = dim_head ** -0.5
self.norm = nn.LayerNorm(dim)
self.attend = nn.Softmax(dim = -1)
self.dropout = nn.Dropout(dropout)
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, dim),
nn.Dropout(dropout)
) if project_out else nn.Identity()
def forward(self, x):
x = self.norm(x)
qkv = self.to_qkv(x).chunk(3, dim = -1)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
attn = self.attend(dots)
attn = self.dropout(attn)
out = torch.matmul(attn, v)
out = rearrange(out, 'b h n d -> b n (h d)')
return self.to_out(out)
class Transformer(nn.Module):
def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(nn.ModuleList([
Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
FeedForward(dim, mlp_dim, dropout = dropout)
]))
def forward(self, x):
for attn, ff in self.layers:
x = attn(x) + x
x = ff(x) + x
return self.norm(x)
class ViT(nn.Module):
def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
super().__init__()
image_height, image_width = pair(image_size)
patch_height, patch_width = pair(patch_size)
assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
num_patches = (image_height // patch_height) * (image_width // patch_width)
patch_dim = channels * patch_height * patch_width
assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
self.to_patch_embedding = nn.Sequential(
Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
nn.LayerNorm(patch_dim),
nn.Linear(patch_dim, dim),
nn.LayerNorm(dim),
)
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
self.dropout = nn.Dropout(emb_dropout)
self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
self.pool = pool
self.to_latent = nn.Identity()
self.mlp_head = nn.Linear(dim, num_classes)
# 解析代码段
def forward(self, img):
x = self.to_patch_embedding(img)
b, n, _ = x.shape
cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = b)
x = torch.cat((cls_tokens, x), dim=1)
x += self.pos_embedding[:, :(n + 1)]
x = self.dropout(x)
x = self.transformer(x)
x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
x = self.to_latent(x)
return self.mlp_head(x)

复制代码

img = torch.randn(1,3,256,256)
b=1
c=3
h=256 = h*p1,h = 8
w =256
self.to_patch.embedding = nn.Sequential(
Rearrange(‘b c (h p1)(w p2)-> b(h w)(p1 p2 c)’,p1 = patch_height, p2 = patch_width),# 图片切分重排
nn.Linear(patch_dim, dim) # Linear Projection of Flattened Patches
)

复制代码

cls_tokens = repeat(self.cls_token, '() n d -> b n d',b = b)
x = torch.cat((cls_tokens, x), dim = 1) # concat方法，维度为1，在n的维度上

复制代码

self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
# 位置编码：它是一个可学习的参数,初始化为随机值。
x += self.pos_embedding[:,:(n+1)]
# 将位置编码加到输入序列上。

复制代码

from __future__ import print_function
import glob
from itertools import chain
import os
import random
import zipfile
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from linformer import Linformer
from PIL import Image
from sklearn.model_selection import train_test_split
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader, DataSet
from torchvision import datasets, transforms
from tqdm.notebook import tqdm
from vit_pytorch import ViT

复制代码

batch_size = 64
epochs = 20
lr = 3e-5
gamma = 0.7
seed = 42

复制代码

def seed_everything(seed):
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
seed_everything(seed)
device = 'cuda'

复制代码

os.makedirs('data',exist_ok = True)
train_dir = 'data/train'
test_dir = 'data/test'
with zipfile.ZipFile('data/train.zip') as train_zip:
train_zip.extractall('data')
with zipfile.ZipFile('data/test.zip') as test_zip:
test_zip.extractall('data')
train_list = glob.glob(os.path.join(train_dir,'*.jpg')) # 查找匹配的jpg文件
test_list = glob.glob(os.path.join(test_dir,'*.jpg'))
print(f'Train Data:{len(train_list)}')
print(f'Test Data:{len(test_list)}')
labels = [path.split('/')[-1].split('\\')[-1].split[0] for path in train_list]
print(train_list[0]
print(labels[0]))

复制代码

random_idx = np.random.randint(1,len(train_list),size = 9)
fig, axes = plt.subplots(3,3,figsize = (16,12))
for idx,ax in enumerate(axes.ravel()):
img = Imag.open(train_list[idx])
ax.set_title(labels[idx])
ax.imshow(img)

复制代码

train_list, valid_list = train_test_split(train_list, test_size = 0.2, stratify = labels, random_state = seed)
print(f'Train Data:{len(train_list)}')
print(f'Validation Data:{len(valid_list)}')
print(f'Test Data:{len(test_list)}')

复制代码

train_tranforms = tranforms.Compose(
[
transforms.Resize((224,224)),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]
)
val_tranforms = tranforms.Compose(
[
transforms.Resize((224,224)),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]
)
test_tranforms = tranforms.Compose(
[
transforms.Resize((224,224)),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]
)

复制代码

class CatsDogsDataset(Dataset):
def __init__(self, file_list, transform = None):
self.file_list = file_list
self.transform = transform
def __len__(self):
self.filelength = len(self.file_list)
return self.filelength
def __getitem__(self,idx):
img_path = self.file_list[idx]
img = Image.open(img_path)
img_transformed = self.transform(img)
label = img_path.split('/')[-1].split("\")[-1].split("、")[0]
label = 1 if label == "dog" else 0
return img_transformed, label

复制代码

train_data = CatsDogsDataset(train_list, transform = train_transforms)
valid_data = CatsDogsDataset(valid_list, transform = val_transforms)
test_data = CatsDogsDataset(test_list, transform = test_transforms)

复制代码

train_loader = DataLoader(dataset = train_data,batch_size = batch_size,shuffle = True)
valid_loader = DataLoader(dataset = valid_data,batch_size = batch_size,shuffle = True)
test_loader = DataLoader(dataset = test_data,batch_size = batch_size,shuffle = True)

复制代码

print(len(train_data), len(train_loader))
print(len(valid_data), len(valid_loader))

复制代码

model = ViT(
image_size = 224,
patch_size = 16,
num_classes = 2,
dim = 768,
depth = 12,
heads = 12,
mlp_dim = 3072,
dropout = 0.1,
emb_dropout = 0.1
).to(device) # 将Transformer模型移动到指定设备上，比如GPU
model.load_state_dict(torch.load('vit_base_patch16_224_r.pth'),strict = False)

复制代码