VITAL Tracker Pytorch 代码阅读笔记
论文链接：https://arxiv.org/pdf/1804.04273.pdf
代码链接：https://github.com/abnerwang/py-Vital
1.项目结构
-
datasets
-
gnet
- g_init.py
- g_pretrain.py
-
models
-
moduels
- model.py
- sample_generator.py
- utils.py
-
pretrain
- data_prov.py
- options_imagenet.yaml
- options_vot.yaml
- prepro_imagenet.py
- prepro_vot.py
- train_mdnet.py
-
tracking
- bbreg.py
- data_prov.py
- gen_config.py
- options.yaml
- run_tracker.py
2.Tracking
刚开始看项目代码无从下手，索性从tracking部分开始吧
import部分：
import numpy as np
import os
import sys
import time
import argparse
import yaml, json
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torch.utils.data as data
import torch.optim as optim
sys.path.insert(0, '.')
from modules.model import MDNet, BCELoss, set_optimizer
from modules.sample_generator import SampleGenerator
from modules.utils import overlap_ratio
from data_prov import RegionExtractor
from bbreg import BBRegressor
from gen_config import gen_config
sys.path.insert(0,'./gnet')#动态加入搜索路径
from gnet.g_init import NetG, set_optimizer_g
from gnet.g_pretrain import *
opts = yaml.safe_load(open('./tracking/options.yaml','r'))
主函数部分：
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-s', '--seq', default='', help='input seq')
parser.add_argument('-j', '--json', default='', help='input json')
parser.add_argument('-f', '--savefig', action='store_true')
parser.add_argument('-d', '--display', action='store_true')
#action的意思是当读取的参数中出现的时候参数字典的建对应的值为True，
args = parser.parse_args()
assert args.seq != '' or args.json != ''
np.random.seed(0)
torch.manual_seed(0)
# Generate sequence config
img_list, init_bbox, gt, savefig_dir, display, result_path = gen_config(args)
这一句对命令行参数进行处理，得到待检测的图像序列img_list，初始的Ground Truth init_bbox, Ground Truth gt，保存图片的目录savefig_dir，是否对结果进行可视化参数display，以及保存结果的地址result_path，下面顺手直接分析gen_config:
import os
import json
import numpy as np
def gen_config(args):
# args的格式是 python tracking/run_tracker.py -s DragonBaby [-d (display fig)] [-f (save fig)]
if args.seq != '':
# generate config from a sequence name
seq_home = '/home/xpwang/datasets/OTB100'
result_home = './results'
seq_name = args.seq
img_dir = os.path.join(seq_home, seq_name, 'img')
gt_path = os.path.join(seq_home, seq_name, 'groundtruth_rect.txt')
img_list = os.listdir(img_dir)
img_list.sort()
img_list = [os.path.join(img_dir, x) for x in img_list]
with open(gt_path) as f:
gt = np.loadtxt((x.replace('\t',',') for x in f), delimiter=',')
init_bbox = gt[0]#初始的框就是gt[0]
result_dir = os.path.join(result_home, seq_name)
if not os.path.exists(result_dir):
os.makedirs(result_dir)
savefig_dir = os.path.join(result_dir, 'figs')
result_path = os.path.join(result_dir, 'result.json')
elif args.json != '':
# load config from a json file
param = json.load(open(args.json, 'r'))
seq_name = param['seq_name']
img_list = param['img_list']
init_bbox = param['init_bbox']
savefig_dir = param['savefig_dir']
result_path = param['result_path']
gt = None
if args.savefig:
if not os.path.exists(savefig_dir):
os.makedirs(savefig_dir)
else:
savefig_dir = ''
return img_list, init_bbox, gt, savefig_dir, args.display, result_path
# Run tracker
result, result_bb, fps = run_vital(img_list, init_bbox, gt=gt, savefig_dir=savefig_dir, display=display)
# Save result
res = {}
res['res'] = result_bb.round().tolist()
res['type'] = 'rect'
res['fps'] = fps
json.dump(res, open(result_path, 'w'), indent=2)
这一步运行跟踪框架并保存结果。
接下来看run_vital()函数。
首先初始化目标框和结果：
def run_vital(img_list, init_bbox, gt=None, savefig_dir='', display=False):
# Init bbox
target_bbox = np.array(init_bbox)
result = np.zeros((len(img_list), 4))
result_bb = np.zeros((len(img_list), 4))
result[0] = target_bbox
result_bb[0] = target_bbox
if gt is not None:
overlap = np.zeros(len(img_list))
overlap[0] = 1
接下来是初始化模型部分：
model = MDNet(opts['model_path'])
model_g = NetG()
if opts['use_gpu']:
model = model.cuda()
model_g = model_g.cuda()
这里是用预训练好的模型加载MDNET，而G网络是初始化生成，看一下MDNET的源码，已经注释：
class MDNet(nn.Module):
def __init__(self, model_path=None, K=1):
super(MDNet, self).__init__()
self.K = K #K是分支数,即一次有多少个视频参与训练，论文中是8个
#搭建网络
self.layers = nn.Sequential(OrderedDict([
('conv1', nn.Sequential(nn.Conv2d(3, 96, kernel_size=7, stride=2),
nn.ReLU(inplace=True),
nn.LocalResponseNorm(2),
nn.MaxPool2d(kernel_size=3, stride=2))),
('conv2', nn.Sequential(nn.Conv2d(96, 256, kernel_size=5, stride=2),
nn.ReLU(inplace=True),
nn.LocalResponseNorm(2),
nn.MaxPool2d(kernel_size=3, stride=2))),
('conv3', nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=1),
nn.ReLU(inplace=True))),
('fc4', nn.Sequential(nn.Linear(512 * 3 * 3, 512),
nn.ReLU(inplace=True))),
('fc5', nn.Sequential(nn.Dropout(0.5),
nn.Linear(512, 512),
nn.ReLU(inplace=True)))]))
#这一步是在fc5后建立多个分支全连接层，每个视频有不同的全连接层
self.branches = nn.ModuleList([nn.Sequential(nn.Dropout(0.5),
nn.Linear(512, 2)) for _ in range(K)])
for m in self.layers.modules():
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0.1)
for m in self.branches.modules():
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
#如果参数已经预训练好，就直接加在
if model_path is not None:
if os.path.splitext(model_path)[1] == '.pth':
self.load_model(model_path)
elif os.path.splitext(model_path)[1] == '.mat':
self.load_mat_model(model_path)
else:
raise RuntimeError('Unkown model format: {:s}'.format(model_path))
self.build_param_dict()
def build_param_dict(self):
self.params = OrderedDict()
for name, module in self.layers.named_children():
# name 为每一层的名称 如：conv1，fc4
# module 为该层的具体网络结构，如：
# Sequential(
# (0): Conv2d(3, 96, kernel_size=(7, 7), stride=(2, 2))
# (1): ReLU()
# (2): LRN()
# (3): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
# )
append_params(self.params, module, name)
append_params(self.params, module, name)
for k, module in enumerate(self.branches):
# k 为分支层的index 0...5
# module同上
append_params(self.params, module, 'fc6_{:d}'.format(k))
def append_params(params, module, prefix):
# params 为参数字典
# module 为具体的模型
# prefix 为模型中具体某一层的名称 如：conv1，fc4，fc6_0
for child in module.children(): # 遍历模型的网络结构
for k,p in child._parameters.items():
if p is None: continue
if isinstance(child, nn.BatchNorm2d):
name = prefix + '_bn_' + k
else:
name = prefix + '_' + k
if name not in params:
params[name] = p
else:
raise RuntimeError('Duplicated param name: {:s}'.format(name))
接着回去看run_tracker部分：
# Init criterion and optimizer
criterion = BCELoss()
criterion_g = torch.nn.MSELoss(reduction='mean')
model.set_learnable_params(opts['ft_layers'])
model_g.set_learnable_params(opts['ft_layers'])
init_optimizer = set_optimizer(model, opts['lr_init'], opts['lr_mult'])
update_optimizer = set_optimizer(model, opts['lr_update'], opts['lr_mult'])
BCELoss是作者新提出的损失函数，细节参考论文3.2：
class BCELoss(nn.Module):
def forward(self, pos_score, neg_score, average=True):
pos_loss = -F.log_softmax(pos_score, dim=1)[:, 1]
pos_loss_p = (torch.ones(pos_loss.size()).cuda() - F.softmax(pos_score, dim=1)[:,1]) * pos_loss
neg_loss = -F.log_softmax(neg_score, dim=1)[:, 0]
neg_loss_p = (torch.ones(neg_loss.size()).cuda() - F.softmax(neg_score, dim=1)[:,0]) * neg_loss
loss = pos_loss_p.sum() + neg_loss_p.sum()
# loss = pos_loss.sum() + neg_loss.sum()
if average:
# loss /= (pos_loss.size(0) + neg_loss.size(0))
loss /= (pos_loss_p.size(0) + neg_loss_p.size(0))
return loss
然后代码设置优化器和可优化的参数，这一步划重点，网络在运行的时候是不对前面的卷积层进行更新的，我们可以去看一下opts[‘ft_layers’]是ft_layers: [‘fc’]，然后再去看set_learnable_params函数：
def set_learnable_params(self, layers):
for k, p in self.params.items():
if any([k.startswith(l) for l in layers]):
p.requires_grad = True
else:
p.requires_grad = False
def get_learnable_params(self):
params = OrderedDict()
for k, p in self.params.items():
if p.requires_grad:
params[k] = p
return params
简而言之，就是fc层都被设置用来更新啦。
然后看设置优化器部分：
def set_optimizer(model, lr_base, lr_mult=opts['lr_mult'], momentum=opts['momentum'], w_decay=opts['w_decay']):
params = model.get_learnable_params() # 得到所有需要求解梯度的参数字典
param_list = []
for k, p in params.items():
lr = lr_base # 设置学习率为0.0001
for l, m in lr_mult.items():
if k.startswith(l):
lr = lr_base * m # fc6层的学习率是0.001
param_list.append({'params': [p], 'lr': lr})
# 随机梯度下降法训练
optimizer = optim.SGD(param_list, lr=lr, momentum=momentum, weight_decay=w_decay)
return optimizer # 返回优化器
好了，都搞清楚了，接着看run_tracker，首先打开第一张图片，然后在bbox生成正负采样：
tic = time.time()
# Load first image
image = Image.open(img_list[0]).convert('RGB')
# Draw pos/neg samples
pos_examples = SampleGenerator('gaussian', image.size, opts['trans_pos'], opts['scale_pos'])(
target_bbox, opts['n_pos_init'], opts['overlap_pos_init'])
neg_examples = np.concatenate([
SampleGenerator('uniform', image.size, opts['trans_neg_init'], opts['scale_neg_init'])(
target_bbox, int(opts['n_neg_init'] * 0.5), opts['overlap_neg_init']),
SampleGenerator('whole', image.size)(
target_bbox, int(opts['n_neg_init'] * 0.5), opts['overlap_neg_init'])])
neg_examples = np.random.permutation(neg_examples)
又要去看函数了，SampleGenrator是一个类，初始化为：
class SampleGenerator():
def __init__(self, type_, img_size, trans=1, scale=1, aspect=None, valid=False):
self.type = type_ # type是指生成哪一种分布的包围盒 'gaussian' 'uniform' 'whole'
self.img_size = np.array(img_size) # (w, h)
self.trans = trans
self.scale = scale
self.aspect = aspect
self.valid = valid
刚刚调用的函数格式是：
pos_examples = SampleGenerator('gaussian', image.size, opts['trans_pos'], opts['scale_pos'])(target_bbox, opts['n_pos_init'], opts['overlap_pos_init'])
去看一下这些参数，把调用翻译一下：
pos_examples = SampleGenerator('gaussian', image.size, 0.1, 1.3)(target_bbox, 500, [0.7,1])
然后再来看下面的函数，overlap_range=[0.7,1], overlap_range 是一个IOU范围,scale_range 指的是缩放范围，接下来请看注释：
def __call__(self, bbox, n, overlap_range=None, scale_range=None):
if overlap_range is None and scale_range is None:
return self._gen_samples(bbox, n)
else:
samples = None
remain = n #初始化，需要生成n个候选框
factor = 2
while remain > 0 and factor < 16:
samples_ = self._gen_samples(bbox, remain * factor)#生成remain*factor个候选框
idx = np.ones(len(samples_), dtype=bool)
if overlap_range is not None: # 挑选出符合IOU范围的候选框
r = overlap_ratio(samples_, bbox)
idx *= (r >= overlap_range[0]) * (r <= overlap_range[1])
if scale_range is not None:
s = np.prod(samples_[:, 2:], axis=1) / np.prod(bbox[2:])
idx *= (s >= scale_range[0]) * (s <= scale_range[1])
samples_ = samples_[idx, :] # 挑选出符合条件的候选框
samples_ = samples_[:min(remain, len(samples_))]
if samples is None:
samples = samples_
else:
samples = np.concatenate([samples, samples_])
remain = n - len(samples)
factor = factor * 2
return samples
下面是刚刚调用的生成样本函数：
def _gen_samples(self, bb, n):
#生成n
# bb: target bbox (min_x,min_y,w,h)
bb = np.array(bb, dtype='float32')
#随便假设一个bbox[1,2,10,5]
# (center_x, center_y, w, h)
sample = np.array([bb[0] + bb[2] / 2, bb[1] + bb[3] / 2, bb[2], bb[3]], dtype='float32')
samples = np.tile(sample[None, :], (n ,1))#纵向复制(center_x, center_y, w, h)
# vary aspect ratio 不同的长宽比
if self.aspect is not None:
ratio = np.random.rand(n, 2) * 2 - 1 # 均匀分布的随机样本[0, 1)×2-1
samples[:, 2:] *= self.aspect ** ratio # w**ratio , h**(-ratio)
# sample generation
if self.type == 'gaussian':
samples[:, :2] += self.trans * np.mean(bb[2:]) * np.clip(0.5 * np.random.randn(n, 2), -1, 1)
samples[:, 2:] *= self.scale ** np.clip(0.5 * np.random.randn(n, 1), -1, 1)
elif self.type == 'uniform':
samples[:, :2] += self.trans * np.mean(bb[2:]) * (np.random.rand(n, 2) * 2 - 1)
samples[:, 2:] *= self.scale ** (np.random.rand(n, 1) * 2 - 1)
elif self.type == 'whole':
m = int(2 * np.sqrt(n))
xy = np.dstack(np.meshgrid(np.linspace(0, 1, m), np.linspace(0, 1, m))).reshape(-1, 2)
xy = np.random.permutation(xy)[:n]
samples[:, :2] = bb[2:] / 2 + xy * (self.img_size - bb[2:] / 2 - 1)
samples[:, 2:] *= self.scale ** (np.random.rand(n, 1) * 2 - 1)
# adjust bbox range
samples[:, 2:] = np.clip(samples[:, 2:], 10, self.img_size - 10)
if self.valid:
samples[:, :2] = np.clip(samples[:, :2], samples[:, 2:] / 2, self.img_size - samples[:, 2:] / 2 - 1)
else:
samples[:, :2] = np.clip(samples[:, :2], 0, self.img_size)
# (min_x, min_y, w, h)
samples[:, :2] -= samples[:, 2:] / 2
return samples
计算bounding box和ground truth的IOU的函数：
def overlap_ratio(rect1, rect2):
'''
Compute overlap ratio between two rects
- rect: 1d array of [x,y,w,h] or
2d array of N x [x,y,w,h]
'''
if rect1.ndim == 1:
rect1 = rect1[None, :]
if rect2.ndim == 1:
rect2 = rect2[None, :]
left = np.maximum(rect1[:, 0], rect2[:, 0])
right = np.minimum(rect1[:, 0] + rect1[:, 2], rect2[:, 0] + rect2[:, 2])
top = np.maximum(rect1[:, 1], rect2[:, 1])
bottom = np.minimum(rect1[:, 1] + rect1[:, 3], rect2[:, 1] + rect2[:, 3])
intersect = np.maximum(0, right - left) * np.maximum(0, bottom - top)
union = rect1[:, 2] * rect1[:, 3] + rect2[:, 2] * rect2[:, 3] - intersect
iou = np.clip(intersect / union, 0, 1)
return iou
OK，继续回到run_vital.py, 往下看，这一步会得到样本框的特征图，换句话说，过VITAL_NET，但只到全连接层之前：
# Extract pos/neg features
pos_feats = forward_samples(model, image, pos_examples)
neg_feats = forward_samples(model, image, neg_examples)
顺便看一下 forward_samples()，我们已经知道这个函数是用来干嘛，再看一下细节吧：
def forward_samples(model, image, samples, out_layer='conv3'):
model.eval()
extractor = RegionExtractor(image, samples, opts)
for i, regions in enumerate(extractor):
if opts['use_gpu']:
regions = regions.cuda()
with torch.no_grad():
feat = model(regions, out_layer=out_layer)
if i==0:
feats = feat.detach().clone()
else:
feats = torch.cat((feats, feat.detach().clone()), 0)
return feats
上面的函数中，RegionExtractor函数用来在图片中取得相应的区域。
class RegionExtractor():
def __init__(self, image, samples, opts):
self.image = np.asarray(image)
self.samples = samples
self.crop_size = opts['img_size']
self.padding = opts['padding']
self.batch_size = opts['batch_test']
self.index = np.arange(len(samples))
self.pointer = 0
def __iter__(self):
return self
def __next__(self):
if self.pointer == len(self.samples): # 采集完了所有的样本区域
self.pointer = 0
raise StopIteration
else:
next_pointer = min(self.pointer + self.batch_size, len(self.samples))
index = self.index[self.pointer:next_pointer]
self.pointer = next_pointer
regions = self.extract_regions(index) # 根据samples提取样本区域
regions = torch.from_numpy(regions)
return regions
next = __next__
def extract_regions(self, index):
regions = np.zeros((len(index), self.crop_size, self.crop_size, 3), dtype='uint8')
for i, sample in enumerate(self.samples[index]): # 根据samples从图像中提取样本区域
regions[i] = crop_image2(self.image, sample, self.crop_size, self.padding)
regions = regions.transpose(0, 3, 1, 2)
regions = regions.astype('float32') - 128.
return regions
从img中选出bbox区域的图片,resize为107*107的from (min_x,min_y,w,h) to (center_x, center_y, w, h):
def crop_image2(img, bbox, img_size=107, padding=16, flip=False, rotate_limit=0, blur_limit=0):
x, y, w, h = np.array(bbox, dtype='float32')
cx, cy = x + w/2, y + h/2
if padding > 0:
w += 2 * padding * w/img_size
h += 2 * padding * h/img_size
# List of transformation matrices
matrices = []
# Translation matrix to move patch center to origin
translation_matrix = np.asarray([[1, 0, -cx],
[0, 1, -cy],
[0, 0, 1]], dtype=np.float32)
matrices.append(translation_matrix)
# Scaling matrix according to image size
scaling_matrix = np.asarray([[img_size / w, 0, 0],
[0, img_size / h, 0],
[0, 0, 1]], dtype=np.float32)
matrices.append(scaling_matrix)
# Define flip matrix
if flip and np.random.binomial(1, 0.5):
flip_matrix = np.eye(3, dtype=np.float32)
flip_matrix[0, 0] = -1
matrices.append(flip_matrix)
# Define rotation matrix
if rotate_limit and np.random.binomial(1, 0.5):
angle = np.random.uniform(-rotate_limit, rotate_limit)
alpha = np.cos(np.deg2rad(angle))
beta = np.sin(np.deg2rad(angle))
rotation_matrix = np.asarray([[alpha, -beta, 0],
[beta, alpha, 0],
[0, 0, 1]], dtype=np.float32)
matrices.append(rotation_matrix)
# Translation matrix to move patch center from origin
revert_t_matrix = np.asarray([[1, 0, img_size / 2],
[0, 1, img_size / 2],
[0, 0, 1]], dtype=np.float32)
matrices.append(revert_t_matrix)
# Aggregate all transformation matrices
matrix = np.eye(3)
for m_ in matrices:
matrix = np.matmul(m_, matrix)
# Warp image, padded value is set to 128
patch = cv2.warpPerspective(img,
matrix,
(img_size, img_size),
borderValue=128)
if blur_limit and np.random.binomial(1, 0.5):
blur_size = np.random.choice(np.arange(1, blur_limit + 1, 2))
patch = cv2.GaussianBlur(patch, (blur_size, blur_size), 0)
return patch
接着看：
# Initial training
train(model, None, criterion, init_optimizer, pos_feats, neg_feats, opts['maxiter_init'])
del init_optimizer, neg_feats
torch.cuda.empty_cache()
g_pretrain(model, model_g, criterion_g, pos_feats)
torch.cuda.empty_cache()
这一步是训练部分，注意，这里能看出来作者论文中的框架了:先更新D，再更新G，接下来分别分析train和g_pretrain：
def g_pretrain(model, model_g, criterion_g, pos_data):
# Evaluate mask
n = pos_data.size(0)
if n % opts['batch_gnet'] == 0:
nBatches = n / opts['batch_gnet']
else:
nBatches = n // opts['batch_gnet'] + 1
print(nBatches)
pos_data = pos_data.view(n, 512, 3, 3)
prob = t.zeros(n)
prob_k = t.zeros(9)
for k in range(9):
row = k // 3
col = k % 3
model.eval()
for i in range(nBatches):
batch = pos_data[opts['batch_gnet'] * i:min(n, opts['batch_gnet'] * (i + 1)), :, :, :]
batch[:, :, col, row] = 0
batch = batch.view(batch.size(0), -1)
if opts['use_gpu']:
batch = batch.cuda()
score = model(batch, in_layer='fc4', out_layer='fc6_softmax')[:, 1]
prob[opts['batch_gnet'] * i:min(n, opts['batch_gnet'] * (i + 1))] = score
model.train()
prob_k[k] = prob.sum()
_, idx = t.min(prob_k, 0)#找到loss最大的mask
idx = idx.item()
row = idx // 3
col = idx % 3
batch_pos = opts['batch_pos']
maxiter_g = opts['maxiter_g']
pos_idx = np.random.permutation(pos_data.size(0))
while len(pos_idx) < batch_pos * maxiter_g:
pos_idx = np.concatenate([pos_idx, np.random.permutation(pos_data.size(0))])
pos_pointer = 0
objective = t.zeros(maxiter_g)
optimizer = set_optimizer_g(model_g)
for iter in range(maxiter_g):
start = time.time()
# select pos idx
pos_next = pos_pointer + batch_pos
pos_cur_idx = pos_idx[pos_pointer:pos_next]
pos_cur_idx = pos_data.new(pos_cur_idx).long()
pos_pointer = pos_next
labels = t.ones(batch_pos, 1, 3, 3)
labels[:, :, col, row] = 0
batch_pos_data = pos_data.index_select(0, pos_cur_idx)
batch_pos_data = batch_pos_data.view(batch_pos_data.size(0), -1)
res = model_g(batch_pos_data)
labels = labels.view(batch_pos, -1)
loss_g = criterion_g(res.float(), labels.cuda().float())
model_g.zero_grad()
loss_g.backward()
optimizer.step()
objective[iter] = loss_g
end = time.time()
print('asdn objective %.3f, %.2f s' % (objective.mean(), end - start))
然后是train部分：
def train(model, model_g, criterion, optimizer, pos_feats, neg_feats, maxiter, in_layer='fc4'):
model.train()
batch_pos = opts['batch_pos']#32
batch_neg = opts['batch_neg']#96
batch_test = opts['batch_test']#256
batch_neg_cand = max(opts['batch_neg_cand'], batch_neg)#1024困难负样本挖掘时生成的1024个负样本，我们要从中取出96个
每次迭代中从样本中随机取出正负样本用来训练
pos_idx = np.random.permutation(pos_feats.size(0)) # 打乱正样本索引
neg_idx = np.random.permutation(neg_feats.size(0))# 打乱负样本索引
while(len(pos_idx) < batch_pos * maxiter):# 每次迭代从正样本中挑选出32个
pos_idx = np.concatenate([pos_idx, np.random.permutation(pos_feats.size(0))])
while(len(neg_idx) < batch_neg_cand * maxiter): # 每次迭代从负样本中挑选出96个
neg_idx = np.concatenate([neg_idx, np.random.permutation(neg_feats.size(0))])
pos_pointer = 0 # 正样本索引
neg_pointer = 0 # 负样本索引
迭代训练maxiter(50)轮：
for i in range(maxiter):
# 挑选正样本idx
pos_next = pos_pointer + batch_pos
pos_cur_idx = pos_idx[pos_pointer:pos_next]
pos_cur_idx = pos_feats.new(pos_cur_idx).long()
pos_pointer = pos_next
# select neg idx
neg_next = neg_pointer + batch_neg_cand
neg_cur_idx = neg_idx[neg_pointer:neg_next]
neg_cur_idx = neg_feats.new(neg_cur_idx).long()
neg_pointer = neg_next
# 创建mini-batch
batch_pos_feats = pos_feats[pos_cur_idx]
if model_g is not None:
batch_asdn_feats = pos_feats.index_select(0, pos_cur_idx)
batch_neg_feats = neg_feats[neg_cur_idx]
# hard negative mining
if batch_neg_cand > batch_neg:
model.eval()
for start in range(0, batch_neg_cand, batch_test):
end = min(start + batch_test, batch_neg_cand)
with torch.no_grad():
score = model(batch_neg_feats[start:end], in_layer=in_layer)
if start==0:
neg_cand_score = score.detach()[:, 1].clone()
else:
neg_cand_score = torch.cat((neg_cand_score, score.detach()[:, 1].clone()), 0)
_, top_idx = neg_cand_score.topk(batch_neg)# 挑选出前96大的，即为困难负样本
batch_neg_feats = batch_neg_feats[top_idx]
model.train()
if model_g is not None:
model_g.eval()
res_asdn = model_g(batch_asdn_feats)#通过g网络，res_asdn的形状是二维向量
model_g.train()
num = res_asdn.size(0)#num=batchsize
mask_asdn = torch.ones(num, 512, 3, 3)#mask的size是n*512*3*3
res_asdn = res_asdn.view(num, 3, 3)#res_asdn的形状是n*3*3
for i in range(num):#循环n次
feat_ = res_asdn[i, :, :]#feat的形状是3*3
featlist = feat_.view(1, 9).squeeze()#一位向量
feat_list = featlist.detach().cpu().numpy()
idlist = feat_list.argsort()#进行排序
idxlist = idlist[:3]#取最小的三个
for k in range(len(idxlist)):
idx = idxlist[k]
row = idx // 3
col = idx % 3
mask_asdn[:, :, col, row] = 0
mask_asdn = mask_asdn.view(mask_asdn.size(0), -1)#mask的size是n*1维向量
if opts['use_gpu']:
batch_asdn_feats = batch_asdn_feats.cuda()
mask_asdn = mask_asdn.cuda()
batch_asdn_feats = batch_asdn_feats * mask_asdn#点乘操作
# forward计算样本得分
if model_g is None:
pos_score = model(batch_pos_feats, in_layer=in_layer)
else:
pos_score = model(batch_asdn_feats, in_layer=in_layer)
neg_score = model(batch_neg_feats, in_layer=in_layer)
# optimize
loss = criterion(pos_score, neg_score)
model.zero_grad()
loss.backward()
if 'grad_clip' in opts:
torch.nn.utils.clip_grad_norm_(model.parameters(), opts['grad_clip'])
optimizer.step()
if model_g is not None:
start = time.time()
prob_k = torch.zeros(9)
for k in range(9):
row = k // 3
col = k % 3
model.eval()
batch = batch_pos_feats.view(batch_pos, 512, 3, 3)
batch[:, :, col, row] = 0
batch = batch.view(batch.size(0), -1)
if opts['use_gpu']:
batch = batch.cuda()
prob = model(batch, in_layer='fc4', out_layer='fc6_softmax')[:, 1]
model.train()
prob_k[k] = prob.sum()
_, idx = torch.min(prob_k, 0)
idx = idx.item()
row = idx // 3
col = idx % 3
optimizer_g = set_optimizer_g(model_g)
labels = torch.ones(batch_pos, 1, 3, 3)
labels[:, :, col, row] = 0
#训练g网络
batch_pos_feats = batch_pos_feats.view(batch_pos_feats.size(0), -1)
res = model_g(batch_pos_feats)
labels = labels.view(batch_pos, -1)
criterion_g = torch.nn.MSELoss(reduction='mean')
loss_g_2 = criterion_g(res.float(), labels.cuda().float())
model_g.zero_grad()
loss_g_2.backward()
optimizer_g.step()
end = time.time()
print('asdn objective %.3f, %.2f s' % (loss_g_2, end - start))
OK，接着往下看，快结束了，再坚持一下：
# 训练bounding box regression
# 生成1000个符合overlab要求和scale要求的bbox
bbreg_examples = gen_samples(SampleGenerator('uniform', image.size, 0.3, 1.5, 1.1),
target_bbox, opts['n_bbreg'], opts['overlap_bbreg'], opts['scale_bbreg'])
bbreg_feats = forward_samples(model, image, bbreg_examples) # conv3的特征，已经转换为fc4的输入格式
bbreg = BBRegressor(image.size) # bounding box regression初始化
bbreg.train(bbreg_feats, bbreg_examples, target_bbox) # bounding box regression训练
这部分也是我没有很仔细看的部分，bounding box regression ,改善定位的精确度，根据给定视频的第一帧，训练一个简单地线性岭回归模型来预测目标物体的位置，用的是 Conv 3 的特征。在随后的视频帧中， 如果预测的目标是可靠的，可以利用bounding box regression调整得到最终的目标位置。文章仅仅用第一帧进行 bounding box regressor 的预测，因为这非常耗时，并且增量学习并不一定有用考虑到其 risk。
# Init sample generators for update
sample_generator = SampleGenerator('gaussian', image.size, opts['trans'], opts['scale'])
pos_generator = SampleGenerator('gaussian', image.size, opts['trans_pos'], opts['scale_pos'])
neg_generator = SampleGenerator('uniform', image.size, opts['trans_neg'], opts['scale_neg'])
# Init pos/neg features for update
# 初始化正负样本特征 for update
neg_examples = neg_generator(target_bbox, opts['n_neg_update'], opts['overlap_neg_init'])
neg_feats = forward_samples(model, image, neg_examples)
pos_feats_all = [pos_feats]
neg_feats_all = [neg_feats]
下面是可视化部分：
spf_total = time.time() - tic
# Display
savefig = savefig_dir != ''
if display or savefig:
dpi = 80.0
figsize = (image.size[0] / dpi, image.size[1] / dpi)
fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.set_axis_off()
fig.add_axes(ax)
im = ax.imshow(image, aspect='auto')
if gt is not None:
gt_rect = plt.Rectangle(tuple(gt[0, :2]), gt[0, 2], gt[0, 3],
linewidth=3, edgecolor="#00ff00", zorder=1, fill=False)
ax.add_patch(gt_rect)
rect = plt.Rectangle(tuple(result_bb[0, :2]), result_bb[0, 2], result_bb[0, 3],
linewidth=3, edgecolor="#ff0000", zorder=1, fill=False)
ax.add_patch(rect)
if display:
plt.pause(.01)
plt.draw()
if savefig:
fig.savefig(os.path.join(savefig_dir, '0000.jpg'), dpi=dpi)
跟论文无关，不赘述了，感觉已经写的太冗长了–
主循环:
for i in range(1, len(img_list)): # 跟踪每一帧
tic = time.time()
# Load image
image = Image.open(img_list[i]).convert('RGB')# 加载图片
# Estimate target bbox
samples = sample_generator(target_bbox, opts['n_samples'])# 每一帧随机生成256个候选区域
sample_scores = forward_samples(model, image, samples, out_layer='fc6')
# 计算这256个候选窗口的得分
top_scores, top_idx = sample_scores[:, 1].topk(5) # 挑选出得分最高的5个候选窗口和索引
top_idx = top_idx.cpu()
target_score = top_scores.mean()# 这一帧的候选框的得分为前5个得分的平均值
target_bbox = samples[top_idx]
if top_idx.shape[0] > 1:
target_bbox = target_bbox.mean(axis=0)# 更新target_bbox用于下一帧的跟踪
success = target_score > 0 # 比较看是否成功
#失败时扩大搜索区域
if success:
sample_generator.set_trans(opts['trans'])
else:
sample_generator.expand_trans(opts['trans_limit'])
# Bbox regression 微调bbox
if success:
bbreg_samples = samples[top_idx]
if top_idx.shape[0] == 1:
bbreg_samples = bbreg_samples[None,:]
bbreg_feats = forward_samples(model, image, bbreg_samples)
bbreg_samples = bbreg.predict(bbreg_feats, bbreg_samples)
bbreg_bbox = bbreg_samples.mean(axis=0)
else:
bbreg_bbox = target_bbox 在失败时复制先前的结果
# Save result
result[i] = target_bbox
result_bb[i] = bbreg_bbox
# Data collect
if success:
# 绘制正负样本
pos_examples = pos_generator(target_bbox, opts['n_pos_update'], opts['overlap_pos_update'])
pos_feats = forward_samples(model, image, pos_examples)
pos_feats_all.append(pos_feats)
if len(pos_feats_all) > opts['n_frames_long']:
del pos_feats_all[0]
neg_examples = neg_generator(target_bbox, opts['n_neg_update'], opts['overlap_neg_update'])
# 提取正负样本的特征
neg_feats = forward_samples(model, image, neg_examples)
neg_feats_all.append(neg_feats)
if len(neg_feats_all) > opts['n_frames_short']:
del neg_feats_all[0]
# Short term update # 跟踪失败时进行短期更新
if not success:
nframes = min(opts['n_frames_short'], len(pos_feats_all))
pos_data = torch.cat(pos_feats_all[-nframes:], 0)
neg_data = torch.cat(neg_feats_all, 0)
train(model, None, criterion, update_optimizer, pos_data, neg_data, opts['maxiter_update'])
# Long term update # 每10帧进行一次长期更新
elif i % opts['long_interval'] == 0:
pos_data = torch.cat(pos_feats_all, 0)
neg_data = torch.cat(neg_feats_all, 0)
train(model, model_g, criterion, update_optimizer, pos_data, neg_data, opts['maxiter_update'])
torch.cuda.empty_cache()
spf = time.time() - tic
spf_total += spf
# Display
if display or savefig:
im.set_data(image)
if gt is not None:
gt_rect.set_xy(gt[i, :2])
gt_rect.set_width(gt[i, 2])
gt_rect.set_height(gt[i, 3])
rect.set_xy(result_bb[i, :2])
rect.set_width(result_bb[i, 2])
rect.set_height(result_bb[i, 3])
if display:
plt.pause(.01)
plt.draw()
if savefig:
fig.savefig(os.path.join(savefig_dir, '{:04d}.jpg'.format(i)), dpi=dpi)
最后的部分：
if gt is None:
print("Frame %d/%d, Score %.3f, Time %.3f" % (i, len(img_list), target_score, spf))
else:
print("Frame %d/%d, Overlap %.3f, Score %.3f, Time %.3f" %
(i, len(img_list), overlap_ratio(gt[i], result_bb[i])[0], target_score, spf))
fps = len(img_list) / spf_total
return result, result_bb, fps
OK，这一部分终于写完了。
3.Pretrain
这一部分实在写不下去了，没什么难度，只是把前面的卷积层一起训练了，划重点，为什么这样呢？因为在offline训练的时候我们训练卷积层就可以获得特征的表达，而后面的全连接层则是分视频来算咯。
Conclusion
花了大概一天的时间搞明白了这个代码，希望对自己和大家都有帮助 ^ ^
python 小知识：
- Orderdict()创建的字典是有序的，按照输入顺序进行输出
- sys.path.insert(0, ‘…/modules’)#动态加入搜索路径
扫一扫
837
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
