【transformer】【ViT】【code】ViT代码

最新推荐文章于 2025-03-09 23:55:53 发布

「已注销」

最新推荐文章于 2025-03-09 23:55:53 发布

阅读量1.2k

点赞数

分类专栏：网络 transformer_CV

本文链接：https://blog.youkuaiyun.com/YJYS_ZHX/article/details/115577594

版权

这篇博客详细介绍了如何实现Vision Transformer (ViT)模型，包括导入einops库进行张量操作，代码调用过程，ViT网络结构以及Block的实现细节。通过阅读，读者可以理解ViT的工作原理并学习其代码实现。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

桃叶儿尖上尖，柳絮儿飞满了天…

1 导入库

import torch
from torch import nn, einsum
import torch.nn.functional as F

from einops import rearrange, repeat
from einops.layers.torch import Rearrange

解释：其中einops库用于张量操作，增强代码的可读性，使用还是比较方便的。教程链接：

einops基础教程
 einsum讲解

2 调用

if __name__=="__main__":
    net = ViT(image_size=256, 
        patch_size=32,#pathces的尺寸
        num_classes=1000,
        dim=1024, #embddings的长度，也就是每个block的输入输出的尺寸
        depth=6,#网络深度，多少个block
        heads=16,#注意力抽头的个数
        mlp_dim=2048,#mlp中反瓶颈结构的中间维度，也就是先升维，再降维
        dropout=0.1,
        emb_dropout=0.1)
    x = torch.rand((2, 3, 256, 256))#测试数据
    output = net(x)

从主干到分支解释代码。

3 ViT网络

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__()
        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
        num_patches = (image_size // patch_size) ** 2
        patch_dim = channels * patch_size ** 2
        assert

最低0.47元/天解锁文章