Pysot训练过程整理（详细注释） - SiamRPN++ & ResNet50
准备阶段
#分布式训练准备：进程号、进程数量初始化。
rank, world_size = dist_init()
#world_size：进程数量/任务数量/GPU数量
#rank：进程ID
#加载参数：合并args参数表（包含代码中定义参数以及experiments中对应config.yaml配置文件）和cfg配置文件。
cfg.merge_from_file(args.cfg)
#merge_from_file(*.yaml)：加载一个yaml配置文件并将其与cfgNode合并
#使用local process（process ID：0）加载log files。
if rank == 0:
if not os.path.exists(cfg.TRAIN.LOG_DIR):
os.makedirs(cfg.TRAIN.LOG_DIR)
init_log('global', logging.INFO)
if cfg.TRAIN.LOG_DIR:
add_file_handler('global',
os.path.join(cfg.TRAIN.LOG_DIR, 'logs.txt'),
logging.INFO)
logger.info("Version Information: \n{}\n".format(commit()))
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
#创建模型：通过ModelBuilder类作为模型各个模块组装类，在构造函数中分别调用了get_backbone/get_neck/get_rpn_head/get_mask_head(可选)
#ModelBuilder实现了训练时用到的前向传播forward(data)函数，输入data包含训练时的template patch和search patch，以及分类label：label_cls、预测框位置回归label：label_loc，以及位置参数权重：label_loc_weight。返回的outputs字典包括了总损失值total_loss/分类损失cls_loss/位置损失loc_loss；
#ModelBuilder类同时实现了推断时模板分支的计算方法template(z)（backbone和neck部分）和搜索分支的计算方法track(x)（backbone和neck部分，以及与模板分支得到的结果一起送入rpn_head部分），得到并返回分类和回归结果cls/loc/mask（可选）。
model = ModelBuilder().cuda().train()
dist_model = DistModule(model)
#model.cuda().train()：将模型移动至GPU上，并set mode=training mode ;
#DistModule(model)：分布式训练模型，需要将模型副本copy至各个GPU中；
下面是ModelBuilder集成的各个模块的实现
1.backbone构建
self.backbone = get_backbone(cfg.BACKBONE.TYPE, **cfg.BACKBONE.KWARGS) #model_builder.py
BACKBONES = {
'alexnetlegacy': alexnetlegacy,
'mobilenetv2': mobilenetv2,
'resnet18': resnet18,
'resnet34': resnet34,
'resnet50': resnet50,
'alexnet': alexnet,
} #__init__.py
def resnet50：model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) #resnet_atrous.py
在这，ResNet50的构建主要由两部分构成：ResNet整体构建和Bottleneck模块构建。
* Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=1)
输入channels=inplanes；输出channels=4*planes；stride=1；dilation=1(default)
输入x => conv1(inplanes, planes) => bn1 => relu => conv2(planes,planes) => bn2 => relu => conv3(planes,planes*4) => bn3 => out+residual(x或x经过downsample层处理后的输出) ，其中conv2可能为空洞卷积。
* ResNet(block, layers, used_layers)
输入图像 => conv1(3,64) => bn1 => x_:relu => maxpool(ksz=3,s=2,p=1) => p1:layer1 => p2:layer2 => p3:layer3 => p4:layer4 => out:[x_,p1,p2,p3,p4] => out[p2,p3,p4]
layer1=_make_layer(block, planes=64, blocks=layers[0], stride=1, dilation=1)：block为Bottleneck实例，layers[0]=3
downsample = nn.Sequential(nn.Conv2d(self.inplanes=64, planes * block.expansion=64*4,kernel_size=1, stride=1, bias=False),nn.BatchNorm2d(planes * block.expansion=64*4),)
block(64,64,1,downsample,1) => Bottleneck(inplanes, planes, stride=1,downsample, dilation=1) => 输入64channels，输出256channels
self.inplanes=64*4=256
block(256,64,dilation=1) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=1) =>输入256channels，输出256channels
block(256,64,dilation=1) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=1) =>输入256channels，输出256channels
返回：nn.Sequential(Bottleneck*3)
功能：下采样、通道升维 64=>256
layer2=_make_layer(block, planes=128, blocks=layers[1], stride=2,dilation=1)：block为Bottleneck实例，layers[1]=4
dd=1,padding=0
downsample = nn.Sequential(nn.Conv2d(self.inplanes=256, planes * block.expansion=128*4,kernel_size=3, stride=2, bias=False,padding=0, dilation=1),nn.BatchNorm2d(planes * block.expansion=128*4),)
block(256,128,2,downsample,1) => Bottleneck(inplanes, planes, stride=2,downsample, dilation=1) => 输入256channels，输出512channels
self.inplanes = 128*4 = 512
block(512,128,dilation=1) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=1) =>输入512channels，输出512channels
block(512,128,dilation=1) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=1) =>输入512channels，输出512channels
block(512,128,dilation=1) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=1) =>输入512channels，输出512channels
返回：nn.Sequential(Bottleneck*4)
功能：下采样、通道升维 256=>512
self.feature_size = 128 * block.expansion = 128*4 = 512
layer3=_make_layer(block, planes=256, blocks=layers[2], stride=1, dilation=2)：block为Bottleneck实例，layers[2]=6
dd = dilation // 2 = 1，padding = dd =1
downsample = nn.Sequential(nn.Conv2d(self.inplanes=512, planes * block.expansion=256*4,kernel_size=3, stride=1, bias=False, padding=1, dilation=1),nn.BatchNorm2d(planes * block.expansion=256*4),)
block(512,256,1,downsample,2) => Bottleneck(inplanes, planes, stride=1,downsample, dilation=2) => 输入512channels，输出1024channels
self.inplanes = 256*4 = 1024
block(1024,256,dilation=2) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=2) =>输入1024channels，输出1024channels
block(1024,256,dilation=2) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=2) =>输入1024channels，输出1024channels
block(1024,256,dilation=2) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=2) =>输入1024channels，输出1024channels
block(1024,256,dilation=2) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=2) =>输入1024channels，输出1024channels
block(1024,256,dilation=2) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=2) =>输入1024channels，输出1024channels
返回：nn.Sequential(Bottleneck*6)
功能：下采样、通道升维 512=>1024、扩大感受野dilation=2
self.feature_size = (256 + 128) * 4 = 1536
layer4=_make_layer(block, planes=512, blocks=layers[3], stride=1, dilation=4)：block为Bottleneck实例，layers[3]=3
dd = dilation // 2 = 2，padding = dd =2
downsample = nn.Sequential(nn.Conv2d(self.inplanes=1024, planes * block.expansion=512*4,kernel_size=3, stride=1, bias=False, padding=2, dilation=2),nn.BatchNorm2d(planes * block.expansion=512*4),)
block(1024,512,1,downsample,4) => Bottleneck(inplanes, planes, stride=1,downsample, dilation=4) => 输入1024channels，输出2048channels
self.inplanes = 512*4 = 2048
block(2048,512,dilation=4) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=4) =>输入2048channels，输出2048channels
block(2048,512,dilation=4) => Bottleneck(inplanes, planes, stride=1,downsample=None, dilation=4) =>输入2048channels，输出2048channels
返回：nn.Sequential(Bottleneck*3)
功能：下采样、通道升维 1024=>2048、扩大感受野dilation=4
self.feature_size = 512 * 4 = 2048
[注意] dilation=1时即为默认的卷积方式，没有空洞；dilation=2表示卷积时cell与cell间空出1个cell。
2.neck构建
if cfg.ADJUST.ADJUST:
self.neck = get_neck(cfg.ADJUST.TYPE, **cfg.ADJUST.KWARGS) #model_builder.py
NECKS = {
'AdjustLayer': AdjustLayer,
'AdjustAllLayer': AdjustAllLayer
} #__init__.py
ADJUST:
ADJUST: true
TYPE: "AdjustAllLayer"
KWARGS:
in_channels: [512, 1024, 2048]
out_channels: [256, 256, 256] #config.yaml
neck层为AdjustAllLayer类型，但其构建也是通过AdjustLayer这个函数实现的。
从config.yaml文件中可以看到，输入/输出channels都是长度为3的列表。因此neck层也会有三个AdjustLayer模块，通过self.add_module()方法添加模块，这三个模块分别称为downsample2/downsample3/downsample4；
看一下AdjustLayer函数的实现：
AdjustLayer(in_channels, out_channels, center_size=7) ，其中in_channels与out_channels每次分别接收来自config.yaml文件中的相应的三个值(512,256)/(1024,256)/(2048,256)
center_size对应论文SiamRPN++中提到的“为了减少计算量，对特征crop center size大小”，center_size=7是默认值，在代码中定义。
AdjustLayer的实现就是一个downsample层：
self.downsample = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
nn.BatchNorm2d(out_channels),
)
训练调用forward函数时，会对经过downsample层处理过后的特征进行判断，是否执行crop：
def forward(self, x):
x = self.downsample(x)
if x.size(3) < 20:
l = (x.size(3) - self.center_size) // 2
r = l + self.center_size
x = x[:, :, l:r, l:r]
return x
3.rpn_head构建
self.rpn_head = get_rpn_head(cfg.RPN.TYPE, **cfg.RPN.KWARGS) #model_builder.py
RPNS = {
'UPChannelRPN': UPChannelRPN,
'DepthwiseRPN': DepthwiseRPN,
'MultiRPN': MultiRPN
} #__init__.py
RPN:
TYPE: 'MultiRPN'
KWARGS:
anchor_num: 5
in_channels: [256, 256, 256]
weighted: true #config.yaml
MultiRPN(anchor_num, in_channels, weighted=False) #rpn.py
这里，RPN的输入均为256channels，SiamRPN++的rpn_head有3个，分别是rpn2/rpn3/rpn4。在经过DepthWiseCorr后，进行加权融合（权值计算：F.softmax(cls_weight)/F.softmax(loc_weight)）；
DepthwiseRPN类分别就分类和回归进行DepthWise-Corr计算：
DepthwiseRPN(anchor_num, in_channels[i], in_channels[i]))
self.cls = DepthwiseXCorr(in_channels, out_channels, 2 * anchor_num)
self.loc = DepthwiseXCorr(in_channels, out_channels, 4 * anchor_num)
其中，in_channels=256, out_channels=256；对应DepthwiseXCorr.init(in_channels=256, hidden=256, out_channels=10(或20), kernel_size=3)
对于模板分支(in:256,out:256)：
self.conv_kernel = nn.Sequential(
nn.Conv2d(in_channels, hidden, kernel_size=kernel_size, bias=False),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
)
对于搜索分支(in:256,out:256)：
self.conv_search = nn.Sequential(
nn.Conv2d(in_channels, hidden, kernel_size=kernel_size, bias=False),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
)
depth-wise correlation：
def xcorr_depthwise(x, kernel):
batch = kernel.size(0)
channel = kernel.size(1)
x = x.view(1, batch*channel, x.size(2), x.size(3))
kernel = kernel.view(batch*channel, 1, kernel.size(2), kernel.size(3))
out = F.conv2d(x, kernel, groups=batch*channel)
out = out.view(batch, channel, out.size(2), out.size(3))
return out
最后通过head层得到输出(in:256,out:10/20)：
self.head = nn.Sequential(
nn.Conv2d(hidden, hidden, kernel_size=1, bias=False),
nn.BatchNorm2d(hidden),
nn.ReLU(inplace=True),
nn.Conv2d(hidden, out_channels, kernel_size=1)
)
最后看一下ModelBuilder的前向传播函数forward(data)：传入的data包含training image pairs、label_cls、label_loc、label_loc_weight****（这些参数key-value pairs会在构建数据集的时候写入）
首先将template/search pairs传入上述构建的模型得到分类结果cls和位置结果loc，然后将cls送入log_softmax得到softmax value：cls；再将cls、label_cls送入select_cross_entropy_loss得到分类损失cls_loss；将loc, label_loc, label_loc_weight送入weight_l1_loss得到回归损失loc_loss。返回ouputs=[ total_loss, cls_loss, loc_loss ]
至此，网络模型构建部分结束。
继续往下说。
#加载backbone预训练权重至刚刚创建的model
if cfg.BACKBONE.PRETRAINED:
cur_path = os.path.dirname(os.path.realpath(__file__))
backbone_path = os.path.join(cur_path, '../', cfg.BACKBONE.PRETRAINED)
load_pretrain(model.backbone, backbone_path)
#插一句：关于加载模型，一般有以下5步：
1.初始化device
device = torch.device('cuda' if cfg.CUDA else 'cpu’) #使用GPU或者CPU
或者 device = torch.cuda.current_device() #直接使用GPU
2.加载模型
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device)) #载入GPU
或者 ckpt = torch.load(args.snapshot, map_location=lambda storage,loc:storage.cpu()) #载入CPU
3.移除前缀（一般用于多GPU分布式训练后的模型）
pretrained_dict = remove_prefix(pretrained_dict['state_dict’], 'module.’) #if "state_dict" in pretrained_dict.keys()
或者 pretrained_dict = remove_prefix(pretrained_dict, 'module.’) #if “state_dict" not in pretrained_dict.keys()
4.检查键完整性：即判断当前构建的模型参数与待加载的模型参数是否匹配
check_keys(model, pretrained_dict)
5.装载参数
model.load_state_dict(pretrained_dict, strict=False)
#构建data loader
#1.构建训练集
train_dataset = TrkDataset()
下面我们看一下数据集是怎么构建的。
首先需要说明的一点是：Dataset类必须要实现索引方法__getitem__方便dataloader获取数据，并要提供数据集的长度length(实现__len__方法)，DataLoader 才能自动地从DataSet中取出小批量的数据集进行训练。
dataset.py => class: TrkDataset => 初始化函数__init__
#初始化函数
def __init__(self,):
super(TrkDataset, self).__init__()
desired_size = (cfg.TRAIN.SEARCH_SIZE - cfg.TRAIN.EXEMPLAR_SIZE) / cfg.ANCHOR.STRIDE + 1 + cfg.TRAIN.BASE_SIZE
if desired_size != cfg.TRAIN.OUTPUT_SIZE:
raise Exception('size not match!')
# create anchor targer（目的：生成anchors）
#1.根据stride/ratios/scales生成anchors（5，4），5:anchor数量; 4:x1,y1,x2,y2.
#2.在25x25的数组中的每一个位置都生成5个anchor_box，每个anchor_box都用(x1,y1,x2,y2)表示，且这里的x1,y1,x2,y2均表示5个anchor_box在25个点映射回原图的坐标
#3.将上一步得到的anchors作为Anchors实例属性all_anchors保存anchor的信息
self.anchor_target = AnchorTarget() # TrkDataset:anchor_target => AnchorTarget:anchors => Anchors:all_anchors
# create sub dataset
self.all_dataset = []
start = 0
self.num = 0
for name in cfg.DATASET.NAMES:
subdata_cfg = getattr(cfg.DATASET, name)
#SubDataset是构建数据集更具体的类，实现的功能有：
#1.加载标签数据train.json；
#2.过滤了标注错误的数据和空目录数据；
#3.将数据重新组装，将同一片段的帧统称为‘frames’：meta_data[video][track]['frames'] = frames
#4.经过筛选后的有效数据作为标签属性:self.labels = meta_data
#5.提供属性：self.num(数据集长度)/self.num_use(实际使用的数据量，当传入值为-1时等于self.num)/self.videos(数据路径列表)/self.pick(随机打乱原有index顺序并截取num_use大小的index列表)/self.start_idx(数据集开始的index)/self.frame_range(隔帧取)
sub_dataset = SubDataset(
name,
subdata_cfg.ROOT,
subdata_cfg.ANNO,
subdata_cfg.FRAME_RANGE,
subdata_cfg.NUM_USE, #VID=10000,others=-1
start #from 0 to +num_VID/+num_COCO/+num_YTB/+num_DET
)
start += sub_dataset.num #每个数据集开始的下标
self.num += sub_dataset.num_use #from 0 to +num_use_VID/+num_use_COCO/+num_use_YTB/+num_use_DET
sub_dataset.log()
self.all_dataset.append(sub_dataset) #self.all_dataset保存加载好的数据集对象
# data augmentation
self.template_aug = Augmentation(
cfg.DATASET.TEMPLATE.SHIFT,
cfg.DATASET.TEMPLATE.SCALE,
cfg.DATASET.TEMPLATE.BLUR,
cfg.DATASET.TEMPLATE.FLIP,
cfg.DATASET.TEMPLATE.COLOR
)
self.search_aug = Augmentation(
cfg.DATASET.SEARCH.SHIFT,
cfg.DATASET.SEARCH.SCALE,
cfg.DATASET.SEARCH.BLUR,
cfg.DATASET.SEARCH.FLIP,
cfg.DATASET.SEARCH.COLOR
)
videos_per_epoch = cfg.DATASET.VIDEOS_PER_EPOCH #600000
self.num = videos_per_epoch if videos_per_epoch > 0 else self.num
self.num *= cfg.TRAIN.EPOCH #60000*20=120,0000
self.pick = self.shuffle() #打乱顺序的index列表
dataset.py => class: TrkDataset => 获取数据方法
def __getitem__(self, index):
index = self.pick[index] #得到即将要读取的数据index
dataset, index = self._find_dataset(index) #返回index所在的数据集以及该index在此数据集中的相对index
#随机对训练数据增广(gray)和增加随机性(neg)
gray = cfg.DATASET.GRAY and cfg.DATASET.GRAY > np.random.random()
neg = cfg.DATASET.NEG and cfg.DATASET.NEG > np.random.random()
# get one dataset
if neg: #通过随机的方式构造patch pair，增加随机性，可认为是negative pairs
# 在dataset中按照index随机取一个，得到 img_path & img_anno
template = dataset.get_random_target(index)
# 在随机一个dataset中选择一个，并随机index得到 img_path & img_anno
search = np.random.choice(self.all_dataset).get_random_target()
else:#选择同一video路径中的image作为patch pair，可认为是positive pairs
template, search = dataset.get_positive_pair(index) #分别得到模板分支和搜索分支的输入img_path & img_anno
# get image
template_image = cv2.imread(template[0])
search_image = cv2.imread(search[0])
# get bounding box (BBox Ground-Truth)
template_box = self._get_bbox(template_image, template[1]) #x1,y1,x2,y2
search_box = self._get_bbox(search_image, search[1])
# augmentation
# template/search均为image
template, _ = self.template_aug(template_image,
template_box,
cfg.TRAIN.EXEMPLAR_SIZE,
gray=gray)
search, bbox = self.search_aug(search_image,
search_box,
cfg.TRAIN.SEARCH_SIZE,
gray=gray)
# get labels
# 调用AnchorTarget类的__call__函数
# __call__函数功能：
# 1.得到分类标签cls
# 2.得到偏移标签delta
# 3.得到偏移权重标签delta_weight
# 4.得到所有anchors与bbox的IOU值overlap
cls, delta, delta_weight, overlap = self.anchor_target(bbox, cfg.TRAIN.OUTPUT_SIZE, neg)
template = template.transpose((2, 0, 1)).astype(np.float32)
search = search.transpose((2, 0, 1)).astype(np.float32)
return {
'template': template,
'search': search,
'label_cls': cls,
'label_loc': delta,
'label_loc_weight': delta_weight,
'bbox': np.array(bbox)
}
#2(可选，分布式).构建数据集切片器
if get_world_size() > 1:
train_sampler = DistributedSampler(train_dataset)
#3.构建data loader
train_loader = DataLoader(train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE,
num_workers=cfg.TRAIN.NUM_WORKERS,
pin_memory=True,
sampler=train_sampler)
#构建优化器和学习率调整器
构建SGD优化器optimizer所需要的参数有：可训练参数列表、动量、权重衰减
构建学习率调整器lr_scheduler所需要的参数有：优化器、训练epochs总数
构建过程分为以下几步：
1.搜集可训练参数写入python列表
trainable_params = []
trainable_params += [{
'params': filter(lambda x: x.requires_grad,model.backbone.parameters()),
'lr': cfg.BACKBONE.LAYERS_LR * cfg.TRAIN.BASE_LR}]
2.构建优化器
optimizer = torch.optim.SGD(trainable_params, momentum=cfg.TRAIN.MOMENTUM, weight_decay=cfg.TRAIN.WEIGHT_DECAY)
3.构建学习率调整器
lr_scheduler = build_lr_scheduler(optimizer, epochs=cfg.TRAIN.EPOCH)
4.更新学习率
lr_scheduler.step(cfg.TRAIN.START_EPOCH)
训练阶段
核心代码：def train(train_loader, model, optimizer, lr_scheduler, tb_writer)
for idx, data in enumerate(train_loader):
if epoch != idx // num_per_epoch + start_epoch:
epoch = idx // num_per_epoch + start_epoch
if get_rank() == 0:
torch.save(
{
'epoch': epoch,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict()},
cfg.TRAIN.SNAPSHOT_DIR+'/checkpoint_e%d.pth' % (epoch))
if epoch == cfg.TRAIN.EPOCH:
return
if cfg.BACKBONE.TRAIN_EPOCH == epoch:
optimizer, lr_scheduler = build_opt_lr(model.module, epoch)
#每一个epoch，更新学习率
lr_scheduler.step(epoch)
cur_lr = lr_scheduler.get_cur_lr()
outputs = model(data)
loss = outputs['total_loss’]
if is_valid_number(loss.data.item()):
optimizer.zero_grad()
loss.backward()
reduce_gradients(model)
if rank == 0 and cfg.TRAIN.LOG_GRADS:
log_grads(model.module, tb_writer, tb_idx)
# clip gradient
clip_grad_norm_(model.parameters(), cfg.TRAIN.GRAD_CLIP)
optimizer.step()
至此，pysot的训练过程梳理完成。训练贯穿了从数据加载到优化、保存模型一整个过程，搞清楚这些任务过程，基本上就搞懂了整个项目。如果有什么地方有疑问，还请大家提出来一起交流，希望对大家有些帮助。