目录
本文是在阅读了TSN论文以及复刻了TSN实验的基础上进行的TSN源码分析。
1.源码准备
直接在github中拉去源代码:
git clone --recursive https://github.com/yjxiong/tsn-pytorch
2.源码结构
源码中的结构如下:
main.py | 训练模型 |
test_models.py | 测试模型 |
opts.py | 参数配置 |
dataset.py | 数据读取 |
models.py | 网络结构构建 |
transforms.py | 数据预处理 |
tf_model_zoo文件夹 | 导入模型结构 |
.tar包 | 运行过后生成的模型 |
本文逐步对上面的文件进行解读。
3.源码分析
3.1数据预处理dataset.py
dataset.py主要是读取数据集,然后通过稀疏采样的方式进行采样。
实现方法主要是定了一个TSNDataSet类,该类继承了torch.utils.DataSet这个基类,通过重写_init_和_getitem_方法来读取数据。
torch.utils.data.Dataset类型的数据并不能作为模型的输入,还要通过torch.utils.data.DataLoader类进一步封装,这是因为数据读取类TSNDataSet返回两个值,第一个值是Tensor类型的数据,第二个值是int型的标签,而torch.utils.data.DataLoader类是将batch size个数据和标签分别封装成一个Tensor,从而组成一个长度为2的list。
3.1.1 初始化__init__
init函数主要是初始化SNDataSet类的参数,源代码如下:
def __init__(self, root_path, list_file,
num_segments=3, new_length=1, modality='RGB',
image_tmpl='img_{:05d}.jpg', transform=None,
force_grayscale=False, random_shift=True, test_mode=False):
self.root_path = root_path
self.list_file = list_file
self.num_segments = num_segments
self.new_length = new_length
self.modality = modality
self.image_tmpl = image_tmpl
self.transform = transform
self.random_shift = random_shift
self.test_mode = test_mode
if self.modality == 'RGBDiff':
self.new_length += 1# Diff needs one more image to calculate diff
self._parse_list()
TSNDataSet类的初始化方法_init_需要如下参数:
root_path : 项目的根目录地址,如果其他文件地址使用绝对地址,则可以写成" "
list_file : 训练或测试的列表文件(.txt文件)地址
num_segments : 视频分割的段数
new_length : 根据输入数据集类型的不同,new_length取不同的值
modality : 输入数据集类型(RGB、光流、RGB差异)
image_tmpl : 图片的名称
transform : 数据集是否进行变换操作
random_shift : 稀疏采样时是否增加一个随机数
test_mode : 是否是测试时的数据集输入
3.1.2 读取图片_load_image
_load_image函数主要是读取图片数据集,源码为
def _load_image(self, directory, idx):
if self.modality == 'RGB' or self.modality == 'RGBDiff':
return [Image.open(os.path.join(directory, self.image_tmpl.format(idx))).convert('RGB')]
elif self.modality == 'Flow':
x_img = Image.open(os.path.join(directory, self.image_tmpl.format('x', idx))).convert('L')
y_img = Image.open(os.path.join(directory, self.image_tmpl.format('y', idx))).convert('L')
return [x_img, y_img]
这个函数根据modality也就是图像种类读取图片,最后用convert函数进行转换(如果不使用.convert('RGB')进行转换的话,读出来的图像是RGBA四通道的,A通道为透明通道,该对深度学习模型训练来说暂时用不到,因此使用convert('RGB')进行通道转换),L表示转化为灰度图像,详细图像类别可以参考博客。
3.1.3 _parse_list函数
_parse_list函数功能在于读取list文件,储存在video_list中
def _parse_list(self):
self.video_list = [VideoRecord(x.strip().split(' ')) for x in open(self.list_file)]
self.video_list是一个长度为训练数据数量的列表。每个值都是VIDEORecord对象,包含一个列表和3个属性,列表长度为3,用空格键分割,分别为帧路径、该视频含有多少帧和帧标签。
3.1.4 _sample_indices函数
_sample_indices函数功能在于实现TSN的稀疏采样,返回的是稀疏采样的帧数列表
def _sample_indices(self, record):
"""
:param record: VideoRecord
:return: list
"""
average_duration = (record.num_frames - self.new_length + 1) // self.num_segments
if average_duration > 0:
offsets = np.multiply(list(range(self.num_segments)), average_duration) + randint(average_duration, size=self.num_segments)
elif record.num_frames > self.num_segments:
offsets = np.sort(randint(record.num_frames - self.new_length + 1, size=self.num_segments))
else:
offsets = np.zeros((self.num_segments,))
return offsets + 1
用一个实际的例子来理解上面代码的作用
假如一个视频有300帧,num_segments=3,输入模态为RGB,稀疏采样的步骤如下:
1.将视频分成num_segments=3段,record.num_frames=150,self.new_length=1,求出平 均 每段的帧数为100帧,即 average_duration=(300+1-1)/3=100
2.定义一个list类型的变量offset,对于第一段,假设随机数 randint(average_duration,size=self.num_segments)=30,第一个片段时 range(self.num_segments)=0,计算可得第一个片段中取到的帧编号为offsets[0]=0*100+30=30
3.同理可获得其他片段中取到帧的编号,假设第二帧时,随机数取40,第三帧时,随机数取20, 计算可得第二个、第三个片段中取到的帧编号,分别为140,220
4.经过上述计算,列表offset=[30,140,220],当返回时,返回的为offset+1,即真正取到的帧数为 [31,141,221]
至于if语句的三种情况就是1.平均每一段都有数据 2.视频帧数小于段数 3.视频帧数为0
3.1.5 _getitem_函数
该函数会在TSNDataSet初始化之后执行,功能在于调用执行稀疏采样的函数_sample_indices或者是_get_val_indices或者_get_test_indices,并且调用get方法,得到TSNDataSet的返回
def __getitem__(self, index):
record = self.video_list[index]
if not self.test_mode:
segment_indices = self._sample_indices(record) if self.random_shift else self._get_val_indices(record)
else:
segment_indices = self._get_test_indices(record)
return self.get(record, segment_indices)
record变量读取的是video_list的第index个数据,包含该视频所在的文件地址、视频包含的帧数和视频所属的分类。if语句判断的是,本次运行是不是测试,如果不是则调用_sample_indices,如果是就调用_get_test_indices。
3.1.6 get函数
get方法的功能在于读取提取的帧图片,并且对帧图片进行变形操作(角裁剪、中心提取等)
def get(self, record, indices):
images = list()
for seg_ind in indices:
p = int(seg_ind)
for i in range(self.new_length):
seg_imgs = self._load_image('/openbayes/home/Tsn/data/rawframes/'+ record.path, p)
images.extend(seg_imgs)
if p < record.num_frames:
p += 1
process_data = self.transform(images)
return process_data, record.label
依次遍历提取到的帧图片,存入images,然后对images里面的图片进行变形。
3.2 模型构建model.py
models.py 的主要功能是对之后的训练模型进行准备。使用一些经典模型作为基础模型,如resnet101、BNInception等,针对不同的输入模态,对最后一层全连接层就行修改,得到我们所需的TSN网络模型。
3.2.1 初始化_init_函数
_init_函数用于对模型的参数初始化。
def __init__(self, num_class, num_segments, modality,
base_model='resnet101', new_length=None,
consensus_type='avg', before_softmax=True,
dropout=0.8,
crop_num=1, partial_bn=True):
super(TSN, self).__init__()
self.modality = modality
self.num_segments = num_segments
self.reshape = True
self.before_softmax = before_softmax
self.dropout = dropout
self.crop_num = crop_num
self.consensus_type = consensus_type
if not before_softmax and consensus_type != 'avg':
raise ValueError("Only avg consensus can be used after Softmax")
if new_length is None:
self.new_length = 1 if modality == "RGB" else 5
else:
self.new_length = new_length
print(("""
Initializing TSN with base model: {}.
TSN Configurations:
input_modality: {}
num_segments: {}
new_length: {}
consensus_module: {}
dropout_ratio: {}
""".format(base_model, self.modality, self.num_segments, self.new_length, consensus_type, self.dropout)))
self._prepare_base_model(base_model)
feature_dim = self._prepare_tsn(num_class)
if self.modality == 'Flow':
print("Converting the ImageNet model to a flow init model")
self.base_model = self._construct_flow_model(self.base_model)
print("Done. Flow model ready...")
elif self.modality == 'RGBDiff':
print("Converting the ImageNet model to RGB+Diff init model")
self.base_model = self._construct_diff_model(self.base_model)
print("Done. RGBDiff model ready.")
self.consensus = ConsensusModule(consensus_type)
if not self.before_softmax:
self.softmax = nn.Softmax()
self._enable_pbn = partial_bn
if partial_bn:
self.partialBN(True)
TSN类的初始化方法_init_需要如下参数:
num_class : 分类后的标签数
num_segments : 视频分割的段数
modality : 输入的模态,RGB、光流、RGB差异
base_model : 基础模型,之后的TSN模型以此基础来修改
new_length : 视频取帧的起点,rgb为1,光流为5
consensus_type : 聚合函数的选择
before_softmax :
dropout : 正则化
crop_num : 数据集修改的类别
partial_bn : 是否部分bn
初始化TSN状态下,基础模型为resnet101、聚合函数选择avg平均池化、正则化0.8、数据集处理类型为1、使用部分批量归一化操作。
另外如果你的输入数据是optical flow或RGBDiff,那么还会对网络结构做修改,分别调用_construct_flow_model方法和_construct_diff_model方法来实现的,主要差别在第一个卷积层,因为该层的输入channel依据不同的输入类型而变化
而后调用**self._prepare_base_model(base_model)和feature_dim = self._prepare_tsn(num_class)**来进一步修改TSN的网络模型结构。
3.2.2 _prepare_base_model
def _prepare_base_model模块主要根据构建网络时选的基础网络如resnet或vgg等,对输入数据集的某些参数进行调整,以满足模型的结构。
if 'resnet' in base_model or 'vgg' in base_model:
self.base_model = getattr(torchvision.models, base_model)(True)
self.base_model.last_layer_name = 'fc'
self.input_size = 224
self.input_mean = [0.485, 0.456, 0.406]
self.input_std = [0.229, 0.224, 0.225]
if self.modality == 'Flow':
self.input_mean = [0.5]
self.input_std = [np.mean(self.input_std)]
elif self.modality == 'RGBDiff':
self.input_mean = [0.485, 0.456, 0.406] + [0] * 3 * self.new_length
self.input_std = self.input_std + [np.mean(self.input_std) * 2] * 3 * self.new_length
以上面代码为例,如果是resnet或者vgg网络为基础模型的话,模型最后一层为fc,输入大小为224,然后做类似归一化操作,参数mean就相当于平移操作,参数std相当于放缩操作。再根据输入图片类型flow或者rgbdiff再进行归一化。
再以BNInception为例
elif base_model == 'BNInception':
import tf_model_zoo
self.base_model = getattr(tf_model_zoo, base_model)()
self.base_model.last_layer_name = 'fc'
self.input_size = 224
self.input_mean = [104, 117, 128]
self.input_std = [1]
if self.modality == 'Flow':
self.input_mean = [128]
elif self.modality == 'RGBDiff':
self.input_mean = self.input_mean * (1 + self.new_length)
首先要从import tf_model_zoo导入配置信息,BNInception类,定义在tf_model_zoo文件夹下的bninception文件夹下的pytorch_load.py中,
manifest = yaml.load(open(model_path))是读进配置好的网络结构(.yml格式),返回的manifest是长度为3的字典,和.yml文件内容对应。其中manifest[‘layers’]是关于网络层的详细定义,其中的每个值表示一个层,每个层也是一个字典,包含数据流关系、名称和结构参数等信息。
3.2.3 _prepare_tsn函数
_prepare_tsn主要是对经过_prepare_base_model函数处理过后的模型进行进一步修改,也可能会对fc层进行调节。
def _prepare_tsn(self, num_class):
feature_dim = getattr(self.base_model, self.base_model.last_layer_name).in_features
if self.dropout == 0:
setattr(self.base_model, self.base_model.last_layer_name, nn.Linear(feature_dim, num_class))
self.new_fc = None
else:
setattr(self.base_model, self.base_model.last_layer_name, nn.Dropout(p=self.dropout))
self.new_fc = nn.Linear(feature_dim, num_class)
std = 0.001
if self.new_fc is None:
normal(getattr(self.base_model, self.base_model.last_layer_name).weight, 0, std)
constant(getattr(self.base_model, self.base_model.last_layer_name).bias, 0)
else:
normal(self.new_fc.weight, 0, std)
constant(self.new_fc.bias, 0)
return feature_dim
(1)用featrue_dim保存最后一层feature_map的通道数
(2)如果要加入dropout层,需要先加入dropout后再加入fc层,否则就直接加入fc(fc层输入是feature_dim,类别为num_class)
(3)然后再是初始化权重和偏置,如果没加入dropout就直接
注:setattr是torch.nn.Module类的一个方法,用来为输入的某个属性赋值,一般可以用来修改网络结构,以setattr(self.base_model, self.base_model.last_layer_name, nn.Dropout(p=self.dropout))为例,输入包含3个值,分别是基础网络,要赋值的属性名,要赋的值。因此当这个setattr语句运行结束后,self.base_model.last_layer_name这一层就是nn.Dropout(p=self.dropout)。
与setattr相对应的getattr方法,就是直接获取网络结构信息。
3.3 main.py训练部分
main.py是训练模型的入口,执行main.py即执行了TSN模型的训练。
3.3.1 AverageMeter类
class AverageMeter(object):
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
AverageMeter类是管理一些变量的更新,比如loss损失、top1准确率等。在初始化的时候,调用reset方法,调用update方法的时候就会进行变量的更新,读取变量直接调用对象属性即可,比如top1准确率top1.val。
3.3.2 adjust_learning_rate函数
adjust_learning_rate是用来调节学习率的
def adjust_learning_rate(optimizer, epoch, lr_steps):
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
decay = 0.1 ** (sum(epoch >= np.array(lr_steps)))
lr = args.lr * decay
decay = args.weight_decay
for param_group in optimizer.param_groups:
param_group['lr'] = lr * param_group['lr_mult']
param_group['weight_decay'] = decay * param_group['decay_mult']
decay保存的是学习率的衰减值,然后更新学习率。
3.3.3 accuracy函数
def accuracy(output, target, topk=(1,)):
"""Computes the precision@k for the specified values of k"""
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].contiguous().view(-1).float().sum(0)
res.append(correct_k.mul_(100.0 / batch_size))
return res
accuracy函数是准确率计算函数,输入output是模型的预测结果,尺寸为batch_size*num_class;target是真实标签,长度为batch_size。maxk为之后要进行排序的前k个数,比如topk=(1,3),则maxk=3,之后找到output中值最大的前三个pred=output.topk(maxk,1,True,True),这一语句调用了pytorch中的topk方法。各参数含义如下:
第一个参数(maxk):需要返回的排序前k个的个数,如maxk=3,则返回前三个较大值
第二个参数(1):表示dim,即按行计算
第三个参数(True):表示largest=True,表示返回的是maxk个最大值
第四个参数(True):表示sorted=True,表示返回排序的结果。
target.view(1,-1).expand_as(pred)将Target的尺寸规范到1batch_size,然后将维度扩充为pred相同的维度,也就是maxkbatchsize,然后调用equal方法计算两个tensor 矩阵相同元素的情况,得到correct是同等维度的矩阵。
correct_k = correct[:k].view(-1).float().sum(0)通过k值来决定计算topk的准确率,sum(0)表示按照列的维度计算和,最后都添加到res列表中返回。
3.3.4 train函数
train函数也就是训练部分,将之前的一些函数封装起来。
def train(train_loader, model, criterion, optimizer, epoch):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
if args.no_partialbn:
model.module.partialBN(False)
else:
model.module.partialBN(True)
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
target = target.cuda(async=True)
input_var = torch.autograd.Variable(input)
target_var = torch.autograd.Variable(target)
# compute output
output = model.forward(input_var)
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1, prec5 = accuracy(output.data, target, topk=(1,5))
losses.update(loss.item(), input.size(0))
top1.update(prec1.item(), input.size(0))
top5.update(prec5.item(), input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
if args.clip_gradient is not None:
total_norm = clip_grad_norm(model.parameters(), args.clip_gradient)
if total_norm > args.clip_gradient:
print("clipping gradient: {} with coef {}".format(total_norm, args.clip_gradient / total_norm))
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
print(('Epoch: [{0}][{1}/{2}], lr: {lr:.5f}\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
epoch, i, len(train_loader), batch_time=batch_time,
data_time=data_time, loss=losses, top1=top1, top5=top5, lr=optimizer.param_groups[-1]['lr'])))
(1) 首先定义了一些变量,这些变量是通过AverageMeter类来管理的。之后判断是否需要部分bn
(2)再调用model.py中的train方法来对模型的参数进行预训练,并且冻结除第一层之外的所有批处理规范化层的均值和方差参数,对全连接层的参数进行训练,达到微调的目的。
(3)然后对train_loader中的数据集进行遍历,存储其中的数据集输入和真实标签。执行output=model(input_var)得到模型的输入结果,调用损失计算函数得到loss,调用之前的accuracy函数来更新top1和top3的准确率。
(4)对于梯度清零、回传和先前学习的过程一样,最后可以得到训练之后的模型和训练的loss、accuracy等。
3.3.5 main函数
main函数主要就是对模型的训练,包含了导入模型、数据准备、训练三个部分。
global args, best_prec1
args = parser.parse_args()
if args.dataset == 'ucf101':
num_class = 8
elif args.dataset == 'hmdb51':
num_class = 51
elif args.dataset == 'kinetics':
num_class = 400
else:
raise ValueError('Unknown dataset '+args.dataset)
model = TSN(num_class, args.num_segments, args.modality,
base_model=args.arch,
consensus_type=args.consensus_type, dropout=args.dropout, partial_bn=not args.no_partialbn)
crop_size = model.crop_size
scale_size = model.scale_size
input_mean = model.input_mean
input_std = model.input_std
policies = model.get_optim_policies()
train_augmentation = model.get_augmentation()
model = torch.nn.DataParallel(model, device_ids=args.gpus).cuda()
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})"
.format(args.evaluate, checkpoint['epoch'])))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
导入模型部分
(1)parser是在opts.py中定义的关于读取命令行参数的对象,保存了模型的一些参数,而后通过判断parser中模型处理数据集的类别来定义本次实验有多少动作类别,本次实验有8个类别。
(2)导入模型,直接调用TSN类,进行初始化,每个参数之前都介绍过(调用过程为初始化TSN类里面的参数,然后调用_prepare_base_model方法导入模型,再调用self._prepare_tsn模型调节模型,最后再根据输入数据集图片类型进行近一步处理,最后这一过程暂时为黑盒)
(3)而后获取模型的变量,再设置多GPU训练模型
(4)args.resume是用来设置是否从断点处开始训练,如果之前没有训练过则不导入,如果之前训练过没有训练完,则可以从之前没有训练完的地方进行训练,或者是直接导入已经训练好的模型。(其中checkpoint=torch.load(args.resume)是用来导入已训练好的模型。model.load_state_dict(checkpoint[‘state_dict’])是完成导入模型的参数初始化。)
数据导入部分
if args.modality != 'RGBDiff':
normalize = GroupNormalize(input_mean, input_std)
else:
normalize = IdentityTransform()
if args.modality == 'RGB':
data_length = 1
elif args.modality in ['Flow', 'RGBDiff']:
data_length = 5
train_loader = torch.utils.data.DataLoader(
TSNDataSet("/openbayes/home/Tsn/data/", "/openbayes/home/Tsn/data/ucf101/ucf101_train_split_1_rawframes.txt", num_segments=args.num_segments,
new_length=data_length,
modality=args.modality,
image_tmpl="img_{:05d}.jpg" if args.modality in ["RGB", "RGBDiff"] else args.flow_prefix+"{}_{:05d}.jpg",
transform=torchvision.transforms.Compose([
train_augmentation,
Stack(roll=args.arch == 'BNInception'),
ToTorchFormatTensor(div=args.arch != 'BNInception'),
normalize,
])),
batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
TSNDataSet("/openbayes/home/data/", "/openbayes/home/data/ucf101/ucf101_val_split_1_rawframes.txt", num_segments=args.num_segments,
new_length=data_length,
modality=args.modality,
image_tmpl="img_{:05d}.jpg" if args.modality in ["RGB", "RGBDiff"] else args.flow_prefix+"{}_{:05d}.jpg",
random_shift=False,
transform=torchvision.transforms.Compose([
GroupScale(int(scale_size)),
GroupCenterCrop(crop_size),
Stack(roll=args.arch == 'BNInception'),
ToTorchFormatTensor(div=args.arch != 'BNInception'),
normalize,
])),
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
(1)根据数据集图片类别有个预处理IdentityTransform和GroupNormalize,都是在transforms.py中的方法,
IdentityTransform方法好像没有进行处理(可能没有理解透彻)
class IdentityTransform(object):
def __call__(self, data):
return data
GroupNormalize方法就是进行了一个减去均值除以标准差的平移放缩操作,获得的这个tensor用normalize保存过后,在数据读取中以他为基准进行归一化处理
class GroupNormalize(object):
def __init__(self, mean, std):
self.mean = mean
self.std = std
def __call__(self, tensor):
rep_mean = self.mean * (tensor.size()[0]//len(self.mean))
rep_std = self.std * (tensor.size()[0]//len(self.std))
# TODO: make efficient
for t, m, s in zip(tensor, rep_mean, rep_std):
t.sub_(m).div_(s)
return tensor
(2)自定义的TSNDataSet类用来处理最原始的数据,返回的是torch.utils.data.Dataset类型,然后通过重写初始化函数_init_和_getitem_方法来读取数据。torch.utils.data.Dataset类型的数据并不能作为模型的输入,还要通过torch.utils.data.DataLoader类进一步封装,将batch_size个数据和标签分别封装成一个Tensor,从而组成一个长度为2的list。
模型训练Train函数调用
数据导入完成后还有优化器选择等操作,再调用train方法或者是validate进行模型训练。
if args.evaluate:
validate(val_loader, model, criterion, 0)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch, args.lr_steps)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
prec1 = validate(val_loader, model, criterion, (epoch + 1) * len(train_loader))
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
Train函数
def train(train_loader, model, criterion, optimizer, epoch):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
if args.no_partialbn:
model.module.partialBN(False)
else:
model.module.partialBN(True)
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
target = target.cuda(async=True)
input_var = torch.autograd.Variable(input)
target_var = torch.autograd.Variable(target)
# compute output
output = model.forward(input_var)
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1, prec5 = accuracy(output.data, target, topk=(1,5))
losses.update(loss.item(), input.size(0))
top1.update(prec1.item(), input.size(0))
top5.update(prec5.item(), input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
if args.clip_gradient is not None:
total_norm = clip_grad_norm(model.parameters(), args.clip_gradient)
if total_norm > args.clip_gradient:
print("clipping gradient: {} with coef {}".format(total_norm, args.clip_gradient / total_norm))
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
变量的更新保存采用自定义的AverageMeter类管理,
数据读取是通过enumerate(train_loader)实现的,是先调用DataLoader类的_iter_方法,该方法里面DataLoaderIter类的_next_方法,在该方法中读取dataset数据时就会调用TSNDataSet类的_getitem_方法,从而完成数据的迭代读取,读取到数据后就将数据从Tensor转换成Variable格式
代码操作:
前向计算:output = model(input_var),得到的output维度就是batch_size*class
损失函数计算: loss = criterion(output, target_var);
准确率计算: prec1, prec5 = accuracy(output.data, target, topk=(1,5))
loss.backward()是损失回传,
optimizer.step()是模型参数更新。
4.函数调用关系图
5.数据流图
6.ipo图及其分析过程
本此分析稀疏采样num_segments=7
6.1过程分析
打印出模型结构(由于太多截取了最开始的一部分以及结尾的一部分)
读取图片经过transform函数过后torch.size如下:
实际上读入模型的数据torch.size如下:
计算input.view((-1,sample_len)+input.size()[-2:])的torch.size:
基础模型也就是resnet101最后一层输出torch.size如下:
微调了fc层后,输出torch.size如下:
self.reshape操作过后troch.size如下:
consensus聚合后维度(将图像预测聚合层视频预测)torch.size如下:
squeeze(1)去掉维度i的torch.size如下:
6.2 过程表格展示
步骤 | torch.size分析 | 备注 |
1.输入ucf数据集 | 图片大小320*240 | width:320,height:240 |
2.dataset初始化,调用transform裁剪图片 | process_data:torch.size[21,224,224] | 21:稀疏采样7张,RGB3*7=21 224:通过transform裁剪操作变成了224*224 |
3.将数据导入model中进行训练时,会执行model(input_var)时,会自动调用model.forward(),函数中会执行base_out=self.base_model(input.view((-1,sample_len)+input.size()[-2:])) | ||
3.1 执行input.size()[-2:]) | torch.size[224,224] | size[-2]表示取最后两个维度 |
3.2执行input.view((-1,sample_len) | torch.size[56,3] | (-1.sample_len)中-1表示根据其他位,自动计算 |
3.3 合并 | torch.size[56,3,224,224] | 直接相加 |
3.4 base_model | base_out: torch.size[56,1024] | base_mode为基础模型,表示模型的最后一层输出的torch.size |
4.微调fc | torch.size[56,8] | |
5 .self.reshape | torch.size[8,7,8] | 片段级预测,7为稀疏采样张数 |
6.consensus | torch.size[8,1,8] | 视频级预测,7张图片为一个视频级 |
7.squeeze(1) | torch.size[8,8] | 去掉维度为1的操作 |
6.3 ipo图