本专栏主要是深度学习/自动驾驶相关的源码实现,获取全套代码请参考

准备

本博客完成Vision Transfomer(VIT)模型的搭建和flowers数据集的训练测试.整个源码包括如下几个任务:
1.读取flowers数据集的dataset类,对应文件dataset.py
2.VIT模型搭建,主要依赖于上几篇博客,对应model.py

1.transfomer中Multi-Head Attention的源码实现的MultiheadAttention类,用于搭建BaseTransformerLayer类,实现encoder和decoder功能
2.transfomer中Decoder和Encoder的base_layer的源码实现的BaseTransformerLayer类,帮助我们丝滑地搭建各类transformer网络
3.transfomer中正余弦位置编码的源码实现[可选]

3.设置优化器学习率和训练/验证模型,对应runner.py和train.py
4.可视化测试单个图片的预测结果,对应demo.py

逐步源码实现

源码结构如下

数据集读取

主要原理:根据dataset的路径,存储各个图片对应的路径,label隐藏在路径中.
在getitem函数中完成指定index图片和label的读取和数据增强功能

class Flowers(Dataset):
    # 用于读取flower数据集
    def __init__(self, dataset_path: str, transforms=None):
        '''
        存储所有数据 data路径和label
        :param dataset_path:
        '''
        super(Dataset, self).__init__()
        flowers = os.listdir(dataset_path)
        flowers = sorted(flowers) # 必须排序,否在每一次顺序不一样训练测试类别就会乱
        self.flower_paths = []
        self.class2label = {}  # 类别str 转 label
        label = 0
        for _, flower in enumerate(flowers):
            flowers_path = os.path.join(dataset_path, flower)
            if os.path.isdir(flowers_path):
                self.class2label[flower] = label
                label +=1
                sub_flowers = os.listdir(flowers_path)
                for sub_flower in sub_flowers:
                    self.flower_paths.append(os.path.join(flowers_path, sub_flower))
        self.label2class = label2class(self.class2label)  # label 转 类别str
        self.transforms = transforms
    ''''''
    def __getitem__(self, item):
        # 读取数据和label
        img = Image.open(self.flower_paths[item])
        label = self.class2label[self.flower_paths[item].split('/')[-2]]
        if self.transforms is not None:
            img = self.transforms(img)  # 数据增强
        return img, label

VIt模型搭建

将整个深度学习模型按照人体分为hand+backbone+neck+head 4个部分,Vit模型不同CNN模型,它的backbone+neck为多个MultiHeadAttention堆叠组成,称之为blocks.

hand

hand主用完成预处理,将数据用"手"揉捏成想要的类型.本处主要完成图片的patch操作,将图片分割成一个个小块,使用大核的卷积完成.然后把w和h拉平后shape就和NLP(b,n,d)一样了.

class PatchLayer(nn.Module):
    def __init__(self, img_size, patch_size=20, embeding_dim=64):
        super(PatchLayer, self).__init__()
        self.grid_size = (img_size[0] // patch_size, img_size[1] // patch_size)
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.proj = nn.Conv2d(in_channels=3,
                              out_channels=embeding_dim,
                              kernel_size=(patch_size, patch_size),
                              stride=patch_size,
                              padding=0)
        self.norm = nn.LayerNorm(normalized_shape=embeding_dim)
    def forward(self, img):
        img = self.proj(img)  # 图片分割
        img = img.flatten(start_dim=2)  # wh拉平
        img = img.permute(0, 2, 1)  # [b wh c]
        img = self.norm(img)
        return img

类别和位置编码

类别编码

直接cat到input上面,那么最后也取出对应的那一列作为类别输出.这是transformer类型网络的常用手段.
个人解释:训练出类别的访问者,这个访问者可以从特征信息(原input)中提取类别信息.训练访问者方法就是类别loss回归,训练时候先果推出因,推理时因推出果

位置编码

add到input上,可以使用可学习式的位置编码也可以使用正余弦位置编码.这是transformer类型网络的常用手段,还要特征层编码等
个人解释:训练出位置的标记者

        # 类别编码
        self.cls_token = nn.Parameter(torch.zeros(size=[1, 1, embed_dim]))
        # 固定位置编码和可学习位置编码
        # self.pos_embed = posemb_sincos_1d(len=num_patches + 1, dim=embed_dim,temperature=1000).unsqueeze(0)
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

blocks

blocks使用注意力机制完成特征提取,
个人解释:
input线性映射为[query,key,value],需求侧(query)从供给侧(value)中取值,取值的根据是qurey@key转置生成的注意力矩阵(需求侧和供给侧每个像素之间的相似度),最后输出与输入shape相同.所以我们重复depth次,多次特征提取.
源码直接调用:transfomer中Decoder和Encoder的base_layer的源码实现的BaseTransformerLayer类

head

主要对transfomer输出的类别特征进行映射,embed维度映射为num_class维度

self.head = nn.Linear(embed_dim, num_classes)

VIT整体

主要是上述几个模块的集合及其正向传播过程:
完成二维图片变一维特征,一维特征transfomer特征提取,分类头输出.

class Vit(nn.Module):
    def __init__(self, img_size=[224, 224], patch_size=16, num_classes=1000,
                 embed_dim=768, depth=12, num_heads=12):
        super(Vit, self).__init__()

        self.patch_embed = PatchLayer(img_size, patch_size, embed_dim)
        num_patches = self.patch_embed.num_patches
        self.blocks = nn.Sequential(*[
            BaseTransformerLayer(attn_cfgs=[dict(embed_dim=embed_dim, num_heads=num_heads)],
                                 fnn_cfg=dict(embed_dim=embed_dim, feedforward_channels=4 * embed_dim, act_cfg='ReLU',
                                              ffn_drop=0.),
                                 operation_order=('self_attn', 'norm', 'ffn', 'norm'))
            for _ in range(depth)
        ])

        # 类别编码
        self.cls_token = nn.Parameter(torch.zeros(size=[1, 1, embed_dim]))
        # 固定位置编码和可学习位置编码
        # self.pos_embed = posemb_sincos_1d(len=num_patches + 1, dim=embed_dim,temperature=1000).unsqueeze(0)
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

        # 分类头
        self.head = nn.Linear(embed_dim, num_classes)

        self.loss_class = nn.CrossEntropyLoss()  # 内置softmax
        self.init_weights()
   ''''''
   def forward(self, img):
        query = self.hand(img)
        query = self.extract_feature(query)
        cls_fea = query[:, -1, :]  # 刚刚class_token被cat到了dim1的最后一个数
        x = self.head(cls_fea)
        return x

Runner(参考mmlab)

建立优化前,设置学习率,根据指定的work_flow顺序进行训练的测试,并保留最优权重

class Runner:
    def __init__(self, arg, model, device):
        self.arg = arg
        # 建立优化器
        params = [p for p in model.parameters() if p.requires_grad]
        self.optimizer = torch.optim.SGD(params=params, lr=arg.lr, momentum=0.9, weight_decay=5E-5)
        lf = lambda x: ((1 + math.cos(x * math.pi / arg.epochs)) / 2) * (1 - arg.lrf) + arg.lrf  # cosine
        self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lf)
        self.model = model.to(device)
        self.device = device
        if arg.load_from is not None and arg.load_from != '':
            weight_dict = torch.load(arg.load_from, map_location=device)
            model.load_state_dict(weight_dict)

    def run(self, dataloaders: dict):
        # 开始训练和验证
        assert 'train' in self.arg.work_flow.keys(), '必须要用训练任务'
        epoch_start = 0
        best_accuracy = 0.0
        while epoch_start < self.arg.epochs:
            for task, times in self.arg.work_flow.items():
                if task == 'train':  # 开始训练
                    for _ in range(times):
                        epoch_start += 1  # epoch只记录训练轮
                        self.model.train()
                        loss_sum = 0.0
                        data_loader = tqdm(dataloaders['train'], file=sys.stdout)
                        for step, data_dict in enumerate(data_loader):
                            img, label = data_dict
                            instance = {
                                'data': img.to(self.device),
                                'label': label.to(self.device)
                            }
                            loss = self.model.loss(**instance)
                            loss_sum += loss.detach()  # 要十分注意 避免往计算图中引入新的东西
                            loss.backward()
                            self.optimizer.step()
                            self.optimizer.zero_grad()
                            data_loader.desc = "[train epoch {}] loss: {:.3f}".\
                                format(epoch_start,loss_sum.item() / (step + 1))
                        self.scheduler.step()
                        print('train: epoch={}, loss={}'.format(epoch_start, loss_sum / (step + 1.0)))

                elif task == 'val':  # 开始验证
                    ''''''
                else:
                    raise ValueError('task must be in [train, val, test]')

可视化

读取单张图片,转换格式输入模型,输出的label,转化为class名和置信度,显示图像,class名和置信度.

if __name__ == '__main__':
    # 建立数据集
    data_transform = transforms.Compose([transforms.Resize(256),
                                         transforms.CenterCrop(224),
                                         transforms.ToTensor(),
                                         transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
    img = Image.open('*****daisy/21652746_cc379e0eea_m.jpg')
    input = data_transform(img).unsqueeze(0)
    label2class = Flowers(dataset_path='../datasets/flower_photos-mini').label2class

    device = torch.device('cuda:0')

    # 建立模型
    model = Vit(img_size=[224, 224],
                patch_size=16,
                embed_dim=768,
                depth=12,
                num_heads=12,
                num_classes=5).to(device)
    weight_dict = torch.load('weights/vit.pth', map_location=device)
    model.load_state_dict(weight_dict)

    model.eval()
    with torch.no_grad():
        output = model(input.to(device))
        output = output.detach().cpu()

        label = output[0].numpy().argmax()
        cnf = torch.softmax(output[0],dim=0).numpy().max()*100.0
        cnf = np.around(cnf, decimals=2) #保留2位小数
    plt.imshow(img)
    plt.title('{} : {}%'.format(label2class[label],cnf))
    plt.show()

总结

vit是视觉transfomer最经典的模型,复现一次代码十分有必要,中间会产生很多思考和问题.
后面章节将会更有价值,我将会:

1.利用本次的代码进行很多思考和trick的验证
2.总结本次代码的BUG们,及其产生的原理和解决方法

如需获取全套代码请参考