Loading... ## 前言 文章链接: https://arxiv.org/abs/2010.11929 源码位置: https://github.com/jeonsworld/ViT-pytorch 视频讲解: https://www.bilibili.com/video/BV15P4y137jb CNN网络利用了归纳偏置信息, 使得可以在中小规模的数据集(imagenet竟然也只算中小规模, 哭死)取得不错的效果. 而Transformer的因为无法利用相关信息, 所以在数据规模有限的情况下表现不如CNN. 但是随着预训练数据集的增大, 到了10M-100M, 开始表现出更大的优势, 因为学习到了其他的特征. 序列构成: 一张224x224x3的图像, 将图像分为16x16x3的连续的小格子, 总数(224/16=14)^2=196个, 等效于构成Transformer中一组句子/序列的全部token, 即此时将图片处理成为了一个序列(长度为196). Embedding的处理: 此时图片的每个token对应图像中16x16x3的区域, 我们借助卷积, 将其转换为长度恰好为16x16x3=768的向量, 卷积核大小16x16, 输入通道3, 输出通道768; PE处理: 图像的每个图像子区域顺序是很重要的, 一旦有变化就不在是同一个图像; 因此位置信息需要带上. 方案有1D/2D多种方案, 最后1D/2D实验对比没有明显差别(分析可能与图像切分的数量不是特别多有关, 1D也可以捕获到足够的信息, 如果是切分成很多的情况, 结果就不一定了). 此处位置处理成了*可学习的参数*, 初始的默认的值均为0 ![ViT结构图](https://zoe.red/usr/uploads/2024/04/3301771983.png) ## 源码 嵌入层与位置编码 * 假定输入图像为224x224x3, 子图PATCH区域为16x16x3(对应一个token) * cls_token: 为了与nlp中Transformer保持一致, 此处设置为全0的向量, 对应的输出在Encoder部分的结果用于经过生成头, 产生最后的分类预测结果 ```python class Embeddings(nn.Module): """Construct the embeddings from patch, position embeddings. """ def __init__(self, config, img_size, in_channels=3): super(Embeddings, self).__init__() self.hybrid = None img_size = _pair(img_size) if config.patches.get("grid") is not None: grid_size = config.patches["grid"] patch_size = (img_size[0] // 16 // grid_size[0], img_size[1] // 16 // grid_size[1]) n_patches = (img_size[0] // 16) * (img_size[1] // 16) self.hybrid = True else: patch_size = _pair(config.patches["size"]) n_patches = (img_size[0] // patch_size[0]) * (img_size[1] // patch_size[1]) # 224/16=14^2=196 self.hybrid = False if self.hybrid: self.hybrid_model = ResNetV2(block_units=config.resnet.num_layers, width_factor=config.resnet.width_factor) in_channels = self.hybrid_model.width * 16 self.patch_embeddings = Conv2d(in_channels=in_channels, out_channels=config.hidden_size, kernel_size=patch_size, stride=patch_size) self.position_embeddings = nn.Parameter(torch.zeros(1, n_patches+1, config.hidden_size)) self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) self.dropout = Dropout(config.transformer["dropout_rate"]) def forward(self, x): B = x.shape[0] cls_tokens = self.cls_token.expand(B, -1, -1) if self.hybrid: x = self.hybrid_model(x) x = self.patch_embeddings(x) x = x.flatten(2) x = x.transpose(-1, -2) x = torch.cat((cls_tokens, x), dim=1) embeddings = x + self.position_embeddings embeddings = self.dropout(embeddings) return embeddings ``` 生成预测结果 * 注意forward部分Encoder输出后部分 ```python class VisionTransformer(nn.Module): def __init__(self, config, img_size=224, num_classes=21843, zero_head=False, vis=False): super(VisionTransformer, self).__init__() self.num_classes = num_classes self.zero_head = zero_head self.classifier = config.classifier self.transformer = Transformer(config, img_size, vis) # 768=16x16*3, 224x224x3的图片切成14x14个 16x16x3的小图 self.head = Linear(config.hidden_size, num_classes) def forward(self, x, labels=None): x, attn_weights = self.transformer(x) # x: [64, 197, 768] [B, img_size/ps^2+1=197,ps*ps*3=768] logits = self.head(x[:, 0]) # 取0号位置的cls_token用于产生成分类; 结果: [64, 10] [B, cls] if labels is not None: loss_fct = CrossEntropyLoss() loss = loss_fct(logits.view(-1, self.num_classes), labels.view(-1)) return loss else: return logits, attn_weights ``` 预训练权重 * 示例: ViT-B_16.npz * 下载页面: [谷歌链接](https://console.cloud.google.com/storage/browser/vit_models/imagenet21k%2Bimagenet2012;tab=objects?prefix=&forceOnObjectsSortingFiltering=false) ![网盘权重位置](https://zoe.red/usr/uploads/2024/04/2337237049.png) 训练代码 * 如果超过显存, 调小`train.py`脚本的训练与评估batch_size ```python python3 train.py --name cifar10-100_500 --dataset cifar10 --model_type ViT-B_16 --pretrained_dir checkpoint/ViT-B_16.npz ``` THE END 本文作者:将夜 本文链接:https://zoe.red/2024/384.html 版权声明:本博客所有文章除特别声明外,均默认采用 CC BY-NC-SA 4.0 许可协议。 最后修改:2024 年 04 月 07 日 © 允许规范转载 赞 如果觉得我的文章对你有用,请随意赞赏