CVPR2021跟踪算法TransT代码详解（Transformer Tracking）

ltr(Learning Tracking Representations)文件

用于学习跟踪表示的基于PyTorch的通用框架。A general PyTorch based framework for learning tracking representations.
安装脚本将自动生成本地配置文件“admin/local.py”。如果未生成该文件，请运行admin.environment.create_default_local_file（）生成该文件。
接下来，设置训练工作区的路径，即保存权重文件的目录。还设置要使用的数据集的路径。
如果所有依赖项都已正确安装，则可以在正确的conda环境中使用run_training.py脚本训练网络。

conda activate pytracking
python run_training.py train_module train_name
1
2

这里，train_module是train_settings中的子模块，train_name是要使用的训练设置文件的名称。

1.admin文件

包括载入网络、tensorboard等功能，还包含环境设置。

1.1 local.py

设置训练工作区的路径，即保存权重文件的目录，还有设置要使用到的数据集在本地的路径。

1.2 settings.py

训练设置，例如数据集和网络路径

1.3 environment.py

配置环境导入数据集

1.4 loading.py

载入网络和模型

1.5 model_constructor.py

装饰器，通过函数构建网络，方便重用代码

1.6 multigpu.py

多GPU训练

1.7 settings.py

训练网络的超参数

1.8 stats.py

计算一些状态值、平均值等等

1.9 tensorboard.py

把训练过程中的日志文件打印出来，方便看训练效果

2.actors文件

包含不同训练的actor类。actor类负责通过网络传递输入数据，从而计算损失。

2.1 base_actor.py

actor的基类，actor类处理数据通过网络的传递并计算损失。

2.1 tracking.py

计算损失

3.dataset文件

包含多个训练数据集的集成，即TrackingNet、GOT-10k、LaSOT、ImageNet VID、DAVIS、YouTube VOS、MS-COCO、SBD、LVIS、ECSD、MSRA10k和HKU-IS。此外，它还包括从图像数据集生成合成视频的模块。

4.data_specs文件

关于不同数据集的训练/验证划分的txt文件。

5.data文件

包含处理数据的功能，例如加载图像、数据扩充、从视频中采样帧。
数据载入，默认使用jpeg4py工具

5.1 processing.py

用于培训TransT的处理类。图像按以下方式处理。
首先，通过添加一些噪音来抖动目标边界框。接下来，从图像中裁剪以抖动目标中心为中心的正方形区域（称为搜索区域），其面积search_area_factor^2倍于抖动框的面积。抖动目标框的原因是避免学习目标总是位于搜索区域中心的偏差。
然后将搜索区域调整为参数search_sz给定的固定大小。
search_area_factor-搜索区域相对于目标大小的大小。
template_area_factor-模板区域相对于模板目标大小的大小。
search_sz-一个整数，表示搜索区域的大小，始终为方形。
temp_sz-一个整数，表示模板区域的大小，始终为方形。
center_jitter_factor-包含在提取搜索区域之前应用于目标中心的抖动量的dict。有关抖动的方式，请参见_get_jittered_box。
scale_jitter_factor-包含提取搜索区域之前应用于目标大小的抖动量的dict。有关抖动的方式，请参见_get_jittered_box。
mode-“pair”或“sequence”。如果mode=“sequence”，则输出具有帧的额外维度。

5.2 sampler.py

采样函数

5.3 transforms.py

对图像进行一些变换，例如加噪音，水平翻转等等。

6.models文件

包含不同层和网络的定义。

6.1 backbone

使用修改版本的ResNet50进行特征提取。具体来说，我们删除了ResNet50的最后一个阶段，并将第四个阶段的输出作为最终输出。我们还将第四阶段的下采样单元的卷积步长从2更改为1，以获得更大的特征分辨率。此外，我们将第四阶段的3×3卷积修改为步长为2的扩张卷积，以增加感受野。

6.2 loss

loss函数

6.3 neck

ECA和CFA 模块

6.4 tracking

transt.py

class TransT(nn.Module):单目标跟踪transt模块

backbone：要使用的backbone的torch模块。参见transt_backline.py
featurefusion_network:featurefusion：特征融合网络的torch模块, transformer的变体。请参见featurefusion_network.py
num_classes：对象类的数量，对于单目标跟踪总是1

   def __init__(self, backbone, featurefusion_network, num_classes):
        super().__init__()
        self.featurefusion_network = featurefusion_network
        hidden_dim = featurefusion_network.d_model
        self.class_embed = MLP(hidden_dim, hidden_dim, num_classes + 1, 3)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size=1)
        self.backbone = backbone
1
2
3
4
5
6
7
8

forward需要一个NestedTensor，它包括：
-search.tensors：批处理图像，形状[批处理大小x 3 x H_search x W_search]
-search.mask：形状为[batch_size x H_search x W_search]的二进制掩码，在填充像素上包含1
-template.tensor:批处理图像，形状[批处理大小x 3 x H_template x W_template]
-emplate.mask：形状为[batch_size x H_template x W_template]的二进制掩码，在填充像素上包含1
它返回包含以下元素的dict：
-“pred_logits”：所有特征向量的分类逻辑。形状=[batch_size x num_vectors x（num_classes+1）]
-“pred_boxes”：所有特征向量的归一化框坐标，表示为（center_x，center_y，高度，宽度）。相对于每个单独图像的大小，这些值在[0，1]中归一化。

   def forward(self, search, template):
        if not isinstance(search, NestedTensor):
            search = nested_tensor_from_tensor(search)
        if not isinstance(template, NestedTensor):
            template = nested_tensor_from_tensor(template)
        feature_search, pos_search = self.backbone(search)
        feature_template, pos_template = self.backbone(template)
        src_search, mask_search= feature_search[-1].decompose()
        assert mask_search is not None
        src_template, mask_template = feature_template[-1].decompose()
        assert mask_template is not None
        hs = self.featurefusion_network(self.input_proj(src_template), mask_template, self.input_proj(src_search), mask_search, pos_template[-1], pos_search[-1])

        outputs_class = self.class_embed(hs)
        outputs_coord = self.bbox_embed(hs).sigmoid()
        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
        return out
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

class SetCriterion(nn.Module):

此类计算TransT的损失。
该过程分为两个步骤：
1） 我们计算真值框和模型输出之间的赋值
2） 我们监督每对匹配的地面真相/预测（监督类和框）

7.train_settings文件

7.1 transt.py

训练过程以及参数设置

    settings.device = 'cuda'
    settings.description = 'TransT with default settings.'
    settings.batch_size = 10 
    settings.num_workers = 4
    settings.multi_gpu = False
    settings.print_interval = 1
    settings.normalize_mean = [0.485, 0.456, 0.406]
    settings.normalize_std = [0.229, 0.224, 0.225]
    settings.search_area_factor = 4.0
    settings.template_area_factor = 2.0
    settings.search_feature_sz = 32
    settings.template_feature_sz = 16
    settings.search_sz = settings.search_feature_sz * 8
    settings.temp_sz = settings.template_feature_sz * 8
    settings.center_jitter_factor = {'search': 3, 'template': 0}
    settings.scale_jitter_factor = {'search': 0.25, 'template': 0}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

Transformer模型的一些参数

    settings.position_embedding = 'sine'
    settings.hidden_dim = 256
    settings.dropout = 0.1
    settings.nheads = 8
    settings.dim_feedforward = 2048
    settings.featurefusion_layers = 4
1
2
3
4
5
6

实例化训练集

   lasot_train = Lasot(settings.env.lasot_dir, split='train')
   got10k_train = Got10k(settings.env.got10k_dir, split='vottrain')
   trackingnet_train = TrackingNet(settings.env.trackingnet_dir, set_ids=list(range(4)))
   coco_train = MSCOCOSeq(settings.env.coco_dir)
1
2
3
4

实例化transform变换

    transform_joint = tfm.Transform(tfm.ToGrayscale(probability=0.05))
    transform_train = tfm.Transform(tfm.ToTensorAndJitter(0.2),
                                    tfm.Normalize(mean=settings.normalize_mean, std=settings.normalize_std))
1
2
3

对训练对进行数据处理

    data_processing_train = processing.TransTProcessing(search_area_factor=settings.search_area_factor,
                                                      template_area_factor = settings.template_area_factor,
                                                      search_sz=settings.search_sz,
                                                      temp_sz=settings.temp_sz,
                                                      center_jitter_factor=settings.center_jitter_factor,
                                                      scale_jitter_factor=settings.scale_jitter_factor,
                                                      mode='sequence',
                                                      transform=transform_train,
                                                      joint_transform=transform_joint)
1
2
3
4
5
6
7
8
9

数据采样（需要用到前面数据处理）

   dataset_train = sampler.TransTSampler([got10k_train], [1],
                                samples_per_epoch=1000*settings.batch_size, max_gap=100, processing=data_processing_train)
                                #lasot_train, got10k_train, coco_train, trackingnet_train, [1,1,1,1]
1
2
3

训练加载器

    loader_train = LTRLoader('train', dataset_train, training=True, batch_size=settings.batch_size, num_workers=settings.num_workers,
                             shuffle=True, drop_last=True, stack_dim=0)
1
2

构造网络

    model = transt_models.transt_resnet50(settings)
1

使用多GPU训练

    if settings.multi_gpu:
        model = MultiGPU(model, dim=0)

    objective = transt_models.transt_loss(settings)
    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('number of params:', n_parameters)

    actor = actors.TranstActor(net=model, objective=objective)
1
2
3
4
5
6
7
8

优化器

    param_dicts = [
        {"params": [p for n, p in model.named_parameters() if "backbone" not in n and p.requires_grad]},
        {
            "params": [p for n, p in model.named_parameters() if "backbone" in n and p.requires_grad],
            "lr": 1e-5,
        },
    ]
    optimizer = torch.optim.AdamW(param_dicts, lr=1e-4,
                                  weight_decay=1e-4)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 500)
1
2
3
4
5
6
7
8
9
10

构造训练器

    trainer = LTRTrainer(actor, [loader_train], optimizer, settings, lr_scheduler)
1

开始训练

    trainer.train(1000, load_latest=True, fail_safe=True)
1

8.run_trainging.py

在train_settings中运行训练脚本。
参数：
train_module：“train_settings/”文件夹中的模块名称，这里使用的是transt。
train_name：设置训练文件的名称。
cudnn_benchmark：是否使用cudnn（默认值为True）。