這篇博客來讀一讀TSN算法的PyTorch代碼，總體而言代碼風格還是不錯的，多讀讀優秀的代碼對自身的提升還是有幫助的，另外因為代碼內容較多，所以分訓練和測試兩篇介紹，這篇介紹訓練代碼，介紹順序為代碼運行順序。TSN算法的介紹可以參考博客TSN(Temporal Segment Networks)算法筆記。
論文：Temporal Segment Networks: Towards Good Practices for Deep Action Recognition
代碼地址：https://github.com/yjxiong/tsn-pytorch

項目結構：
main.py是訓練腳本
test_models.py是測試腳本
opts.py是參數配置腳本
dataset.py是數據讀取腳本
models.py是網絡結構構建腳本
transforms.py是數據預處理相關的腳本
tf_model_zoo文件夾關于導入模型結構的腳本

main.py是訓練模型的入口。
首先是導入模塊，其中比較重要的是導入模型：from models import TSN，導入配置的參數：from opts import parser。

import argparse import os import time import shutil import torch import torchvision import torch.nn.parallel import torch.backends.cudnn as cudnn import torch.optim from torch.nn.utils import clip_grad_normfrom dataset import TSNDataSet from models import TSN from transforms import * from opts import parserbest_prec1 = 0

main函數主要包含導入模型、數據準備、訓練三個部分，接下來將按順序介紹。parser是在opts.py中定義的關于讀取命令行參數的對象，然后通過from opts import parser導入的。model = TSN(num_class, args.num_segments, args.modality,...，partial_bn=not args.no_partialbn)這一行是導入模型操作，TSN類的定義在models.py腳本中。輸入包含分類的類別數：num_class；args.num_segments表示把一個video分成多少份，對應論文中的K，默認K=3；采用哪種輸入：args.modality，比如RGB表示常規圖像，Flow表示optical flow等；采用哪種模型：args.arch，比如resnet101，BNInception等；不同輸入snippet的融合方式：args.consensus_type，比如avg等；dropout參數：args.dropout。

def main():global args, best_prec1args = parser.parse_args()if args.dataset == 'ucf101':num_class = 101elif args.dataset == 'hmdb51':num_class = 51elif args.dataset == 'kinetics':num_class = 400else: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)

TSN類（定義在models.py中）的初始化操作：__init__，這里只列出主要的代碼。new_length和輸入數據類型相關。這里主要調用了該類的兩個方法來完成初始化操作，一個是self._prepare_base_model(base_model)，通過調用TSN類的_prepare_base_model方法來導入模型。另一個是feature_dim = self._prepare_tsn(num_class)，通過調用TSN類的_prepare_tsn方法來得到。另外如果你的輸入數據是optical flow或RGBDiff，那么還會對網絡結構做修改，分別調用_construct_flow_model方法和_construct_diff_model方法來實現的，主要差別在第一個卷積層，因為該層的輸入channel依據不同的輸入類型而變化。接下來依次介紹這些方法。

class TSN(nn.Module):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__()if new_length is None:self.new_length = 1 if modality == "RGB" else 5else:self.new_length = new_lengthself._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_bnif partial_bn:self.partialBN(True)

_prepare_base_model方法的部分代碼（以base_model為‘BNInception為例’）如下。getattr模塊的使用：getattr(tf_model_zoo, base_model)()類似tf_model_zoo.BNInception()，因為要根據base_model的不同指定值來導入不同的網絡，所以用getattr模塊。導入模型之后就是一些常規的配置信息了。

elif base_model == 'BNInception':import tf_model_zooself.base_model = getattr(tf_model_zoo, base_model)()self.base_model.last_layer_name = 'fc'self.input_size = 224self.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)

BNInception類，定義在tf_model_zoo文件夾下的bninception文件夾下的pytorch_load.py中。前面當運行self.base_model = getattr(tf_model_zoo, base_model)()，且base_model是‘BNInception’的時候就會調用這個BNInception類的初始化函數__init__。manifest = yaml.load(open(model_path))是讀進配置好的網絡結構（.yml格式），返回的manifest是長度為3的字典，和.yml文件內容對應。其中manifest[‘layers’]是關于網絡層的詳細定義，其中的每個值表示一個層，每個層也是一個字典，包含數據流關系、名稱和結構參數等信息。然后get_basic_layer函數是用來根據這些參數得到具體的網絡層并保存相關信息。setattr(self, id, module)是將得到的層寫入self的指定屬性中，就是搭建層的過程。這樣循環完所有層的配置信息后，就搭建好了整個網絡。
構建好了網絡結構后，另外比較重要的是：self.load_state_dict(torch.utils.model_zoo.load_url(weight_url))這一行，可以分解一下，里面的torch.utils.model_zoo.load_url(weight_url)是通過提供的.pth文件的url地址來下載指定的.pth文件，在PyTorch中.pth文件就是模型的參數文件，如果你已經有合適的模型了且不想下載，那么可以通過torch.load(‘the/path/of/.pth’)導入，因為torch.utils.model_zoo.load_url方法最后返回的時候也是用torch.load接口封裝成字典輸出。self.load_state_dict()則是將導入的模型參數賦值到self中。因此不想下載的話可以用checkpoint=torch.load('the/path/of/.pth')和self.load_state_dict(checkpoint)兩行代替self.load_state_dict(torch.utils.model_zoo.load_url(weight_url))。

class BNInception(nn.Module):def __init__(self, model_path='tf_model_zoo/bninception/bn_inception.yaml', num_classes=101,weight_url='https://yjxiong.blob.core.windows.net/models/bn_inception-9f5701afb96c8044.pth'):super(BNInception, self).__init__()manifest = yaml.load(open(model_path))layers = manifest['layers']self._channel_dict = dict()self._op_list = list()for l in layers:out_var, op, in_var = parse_expr(l['expr'])if op != 'Concat':id, out_name, module, out_channel, in_name = get_basic_layer(l,3 if len(self._channel_dict) == 0 else self._channel_dict[in_var[0]],conv_bias=True)self._channel_dict[out_name] = out_channelsetattr(self, id, module)self._op_list.append((id, op, out_name, in_name))else:self._op_list.append((id, op, out_var[0], in_var))channel = sum([self._channel_dict[x] for x in in_var])self._channel_dict[out_var[0]] = channelself.load_state_dict(torch.utils.model_zoo.load_url(weight_url))

_prepare_tsn方法。feature_dim是網絡最后一層的輸入feature map的channel數。接下來如果有dropout層，那么添加一個dropout層后連一個全連接層，否則就直接連一個全連接層。setattr是torch.nn.Module類的一個方法，用來為輸入的某個屬性賦值，一般可以用來修改網絡結構，以setattr(self.base_model, self.base_model.last_layer_name, nn.Dropout(p=self.dropout))為例，輸入包含3個值，分別是基礎網絡，要賦值的屬性名，要賦的值，一般而言setattr的用法都是這樣。因此當這個setattr語句運行結束后，self.base_model.last_layer_name這一層就是nn.Dropout(p=self.dropout)。
最后對全連接層的參數（weight）做一個0均值且指定標準差的初始化操作，參數（bias）初始化為0。getattr同樣是torch.nn.Module類的一個方法，與為屬性賦值方法setattr相比，getattr是獲得屬性值，一般可以用來獲取網絡結構相關的信息，以getattr(self.base_model, self.base_model.last_layer_name)為例，輸入包含2個值，分別是基礎網絡和要獲取值的屬性名。

def _prepare_tsn(self, num_class):feature_dim = getattr(self.base_model, self.base_model.last_layer_name).in_featuresif self.dropout == 0:setattr(self.base_model, self.base_model.last_layer_name, nn.Linear(feature_dim, num_class))self.new_fc = Noneelse: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.001if 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

前面提到如果輸入不是RGB，那么就要修改網絡結構，這里以models.py腳本中TSN類的_construct_flow_model方法介紹對于optical flow類型的輸入需要修改哪些網絡結構。conv_layer是第一個卷積層的內容，params 包含weight和bias，kernel_size就是(64,3,7,7)，因為對于optical flow的輸入，self.new_length設置為5，所以new_kernel_size是(63,10,7,7)。new_kernels是修改channel后的卷積核參數，主要是將原來的卷積核參數復制到新的卷積核。然后通過nn.Conv2d來重新構建卷積層。new_conv.weight.data = new_kernels是賦值過程。

def _construct_flow_model(self, base_model):# modify the convolution layers# Torch models are usually defined in a hierarchical way.# nn.modules.children() return all sub modules in a DFS mannermodules = list(self.base_model.modules())first_conv_idx = list(filter(lambda x: isinstance(modules[x], nn.Conv2d), list(range(len(modules)))))[0]conv_layer = modules[first_conv_idx]container = modules[first_conv_idx - 1]# modify parameters, assume the first blob contains the convolution kernelsparams = [x.clone() for x in conv_layer.parameters()]kernel_size = params[0].size()new_kernel_size = kernel_size[:1] + (2 * self.new_length, ) + kernel_size[2:]new_kernels = params[0].data.mean(dim=1, keepdim=True).expand(new_kernel_size).contiguous()new_conv = nn.Conv2d(2 * self.new_length, conv_layer.out_channels,conv_layer.kernel_size, conv_layer.stride, conv_layer.padding,bias=True if len(params) == 2 else False)new_conv.weight.data = new_kernelsif len(params) == 2:new_conv.bias.data = params[1].data # add bias if neccessarylayer_name = list(container.state_dict().keys())[0][:-7] # remove .weight suffix to get the layer name# replace the first convlution layersetattr(container, layer_name, new_conv)return base_model

接著main函數的思路，前面這幾行都是在TSN類中定義的變量或者方法，model = torch.nn.DataParallel(model, device_ids=args.gpus).cuda()是設置多GPU訓練模型。args.resume這個參數主要是用來設置是否從斷點處繼續訓練，比如原來訓練模型訓到一半停止了，希望繼續從保存的最新epoch開始訓練，因此args.resume要么是默認的None，要么就是你保存的模型文件（.pth）的路徑。其中checkpoint = torch.load(args.resume)是用來導入已訓練好的模型。model.load_state_dict(checkpoint[‘state_dict’])是完成導入模型的參數初始化model這個網絡的過程，load_state_dict是torch.nn.Module類中重要的方法之一。

crop_size = model.crop_sizescale_size = model.scale_sizeinput_mean = model.input_meaninput_std = model.input_stdpolicies = 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)))cudnn.benchmark = True

介紹完第一部分模型導入后，接下來是main函數中的第二部分：數據導入。首先是自定義的TSNDataSet類用來處理最原始的數據，返回的是torch.utils.data.Dataset類型，一般而言在PyTorch中自定義的數據讀取類都要繼承torch.utils.data.Dataset這個基類，比如此處的TSNDataSet類，然后通過重寫初始化函數__init__和__getitem__方法來讀取數據。torch.utils.data.Dataset類型的數據并不能作為模型的輸入，還要通過torch.utils.data.DataLoader類進一步封裝，這是因為數據讀取類TSNDataSet返回兩個值，第一個值是Tensor類型的數據，第二個值是int型的標簽，而torch.utils.data.DataLoader類是將batch size個數據和標簽分別封裝成一個Tensor，從而組成一個長度為2的list。對于torch.utils.data.DataLoader類而言，最重要的輸入就是TSNDataSet類的初始化結果，其他如batch size和shuffle參數是常用的。通過這兩個類讀取和封裝數據，后續再轉為Variable就能作為模型的輸入了。

# Data loading codeif args.modality != 'RGBDiff':normalize = GroupNormalize(input_mean, input_std)else:normalize = IdentityTransform()if args.modality == 'RGB':data_length = 1elif args.modality in ['Flow', 'RGBDiff']:data_length = 5train_loader = torch.utils.data.DataLoader(TSNDataSet("", args.train_list, 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("", args.val_list, 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)

自定義數據讀取相關類的時候需要繼承torch.utils.data.Dataset這個基類。在TSNDataSet類的初始化函數__init__中最重要的是self._parse_list()，也就是調用了該類的_parse_list()方法。在該方法中，self.list_file就是訓練或測試的列表文件（.txt文件），里面包含三列內容，用空格鍵分隔，第一列是video名，第二列是video的幀數，第三列是video的標簽。VideoRecord這個類只是提供了一些簡單的封裝，用來返回關于數據的一些信息（比如幀路徑、該視頻包含多少幀、幀標簽）。因此最后self.video_list的內容就是一個長度為訓練數據數量的列表，列表中的每個值都是VideoRecord對象，該對象包含一個列表和3個屬性，列表長度為3，分別是幀路徑、該視頻包含多少幀、幀標簽，同樣這三者也是三個屬性的值。

class TSNDataSet(data.Dataset):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_pathself.list_file = list_fileself.num_segments = num_segmentsself.new_length = new_lengthself.modality = modalityself.image_tmpl = image_tmplself.transform = transformself.random_shift = random_shiftself.test_mode = test_modeif self.modality == 'RGBDiff':self.new_length += 1# Diff needs one more image to calculate diffself._parse_list()def _parse_list(self):self.video_list = [VideoRecord(x.strip().split(' ')) for x in open(self.list_file)]

介紹完第二部分數據讀取后，接下來就是main函數的第三部分：訓練模型。這里包括定義損失函數、優化函數、一些超參數設置等，然后訓練模型并在指定epoch驗證和保存模型。adjust_learning_rate(optimizer, epoch, args.lr_steps)是設置學習率變化策略，args.lr_steps是一個列表，里面的值表示到達多少個epoch的時候要改變學習率，在adjust_learning_rate函數中，默認是修改學習率的時候修改成當前的0.1倍。train(train_loader, model, criterion, optimizer, epoch)就是訓練模型，輸入包含訓練數據、模型、損失函數、優化函數和要訓練多少個epoch。最后的if語句是當訓練epoch到達指定值的時候就進行一次模型驗證和模型保存，args.eval_freq這個參數就是用來控制保存的epoch值。prec1 = validate(val_loader, model, criterion, (epoch + 1) * len(train_loader))就是用訓練好的模型驗證測試數據集。最后的save_checkpoint函數就是保存模型參數（model）和其他一些信息，這里我對源代碼做了修改，希望有助于理解，該函數中主要就是調用torch.save(mode, save_path)來保存模型。模型訓練函數train和模型驗證函數validate函數是重點，后面詳細介紹。

# define loss function (criterion) and optimizerif args.loss_type == 'nll':criterion = torch.nn.CrossEntropyLoss().cuda()else:raise ValueError("Unknown loss type")for group in policies:print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(group['name'], len(group['params']), group['lr_mult'], group['decay_mult'])))optimizer = torch.optim.SGD(policies,args.lr,momentum=args.momentum,weight_decay=args.weight_decay)if args.evaluate:validate(val_loader, model, criterion, 0)returnfor epoch in range(args.start_epoch, args.epochs):adjust_learning_rate(optimizer, epoch, args.lr_steps)# train for one epochtrain(train_loader, model, criterion, optimizer, epoch)# evaluate on validation setif (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 checkpointis_best = prec1 > best_prec1best_prec1 = max(prec1, best_prec1)save_checkpoint(epoch=epoch + 1, arch=args.arch, state_dict=model, is_best=is_best)def save_checkpoint(epoch, arch, model, is_best):filename = os.path.join(args.snapshot_pref, '_'.join((args.modality.lower(), 'arch:{}', 'epoch:{}', 'checkpoint.pth')).format(arch, epoch))torch.save(model, filename)if is_best:best_name = os.path.join(args.snapshot_pref, '_'.join((args.modality.lower(), 'arch:{}', 'epoch:{}', 'model_best.pth')).format(arch, epoch))shutil.copyfile(filename, best_name)

train函數是模型訓練的入口。首先一些變量的更新采用自定義的AverageMeter類來管理，后面會介紹該類的定義。然后model.train()是設置為訓練模式。 for i, (input, target) in enumerate(train_loader) 是數據迭代讀取的循環函數，具體而言，當執行enumerate(train_loader)的時候，是先調用DataLoader類的__iter__方法，該方法里面再調用DataLoaderIter類的初始化操作__init__。而當執行for循環操作時，調用DataLoaderIter類的__next__方法，在該方法中通過self.collate_fn接口讀取self.dataset數據時就會調用TSNDataSet類的__getitem__方法，從而完成數據的迭代讀取。讀取到數據后就將數據從Tensor轉換成Variable格式，然后執行模型的前向計算：output = model(input_var)，得到的output就是batch size*class維度的Variable；損失函數計算： loss = criterion(output, target_var)；準確率計算： prec1, prec5 = accuracy(output.data, target, topk=(1,5))；模型參數更新等等。其中loss.backward()是損失回傳， optimizer.step()是模型參數更新。

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 modemodel.train()end = time.time()for i, (input, target) in enumerate(train_loader):# measure data loading timedata_time.update(time.time() - end)target = target.cuda(async=True)input_var = torch.autograd.Variable(input)target_var = torch.autograd.Variable(target)# compute outputoutput = model(input_var)loss = criterion(output, target_var)# measure accuracy and record lossprec1, prec5 = accuracy(output.data, target, topk=(1,5))losses.update(loss.data[0], input.size(0))top1.update(prec1[0], input.size(0))top5.update(prec5[0], input.size(0))# compute gradient and do SGD stepoptimizer.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 timebatch_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'])))

前面提到在train函數中采用自定義的AverageMeter類來管理一些變量的更新。在初始化的時候就調用的重置方法reset。當調用該類對象的update方法的時候就會進行變量更新，當要讀取某個變量的時候，可以通過對象.屬性的方式來讀取，比如在train函數中的top1.val讀取top1準確率。

class AverageMeter(object):"""Computes and stores the average and current value"""def __init__(self):self.reset()def reset(self):self.val = 0self.avg = 0self.sum = 0self.count = 0def update(self, val, n=1):self.val = valself.sum += val * nself.count += nself.avg = self.sum / self.count

前面提到在運行for i, (input, target) in enumerate(train_loader)的時候最終會調用TSNDataSet類的__getitem__方法，該方法就是用來返回具體數據的。前面介紹過TSNDataSet類的初始化函數__init__，在那里面都是一些初始化或定義操作，真正的數據讀取操作是在__getitem__方法中。在__getitem__方法中，record = self.video_list[index]得到的record就是一幀圖像的信息，index是隨機的，這個和前面數據讀取中的shuffle參數對應。在訓練的時候，self.test_mode是False，所以執行if語句，另外self.random_shift默認是True，所以最后執行的是segment_indices = self._sample_indices(record)。在測試的時候，會設置self.test_mode為True，這樣的話就會執行segment_indices = self._get_test_indices(record)。最后再通過get方法返回。接下來分別介紹這三個方法。

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)

在TSNDataSet類的_sample_indices方法中，average_duration表示某個視頻分成self.num_segments份的時候每一份包含多少幀圖像，因此只要該視頻的總幀數大于等于self.num_segments，就會執行if average_duration > 0這個條件，在該條件語句下offsets的計算分成兩部分，np.multiply(list(range(self.num_segments)), average_duration)相當于確定了self.num_segments個片段的區間，randint(average_duration, size=self.num_segments)則是生成了self.num_segments個范圍在0到average_duration的數值，二者相加就相當于在這self.num_segments個片段中分別隨機選擇了一幀圖像。因此在__getitem__方法中返回的segment_indices就是一個長度為self.num_segments的列表，表示幀的index。

def _sample_indices(self, record):""":param record: VideoRecord:return: list"""average_duration = (record.num_frames - self.new_length + 1) // self.num_segmentsif 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

在TSNDataSet類的_get_test_indices方法中，就是將輸入video按照相等幀數距離分成self.num_segments份，最終返回的offsets就是長度為self.num_segments的numpy array，表示從輸入video中取哪些幀作為模型的輸入。該方法是模型測試的時候才會調用。

def _get_test_indices(self, record):tick = (record.num_frames - self.new_length + 1) / float(self.num_segments)offsets = np.array([int(tick / 2.0 + tick * x) for x in range(self.num_segments)])return offsets + 1

在TSNDataSet類的get方法中，先通過seg_imgs = self._load_image(record.path, p)來讀取圖像數據。_load_image方法中主要就是采用PIL庫的Image模塊來讀取圖像數據，該方法比較固定，一般作為當前類的一個方法比較合適，另外區分RGB和Flow數據讀取的原因主要是圖像名稱不同。對于RGB或RGBDiff數據，返回的seg_imgs是一個長度為1的列表，對于Flow數據，返回的seg_imgs是一個長度為2的列表，然后將讀取到的圖像數據合并到images這個列表中。另外對于RGB而言，self.new_length是1，這樣images的長度就是indices的長度；對于Flow而言，self.new_length是5，這樣images的長度就是indices的長度乘以(5*2)。process_data = self.transform(images)將list類型的images封裝成了Tensor，在訓練的時候：對于RGB輸入，這個Tensor的尺寸是(3*self.num_segments,224,224)，其中3表示3通道彩色；對于Flow輸入，這個Tensor的尺寸是(self.num_segments*2*self.new_length,224,224)，其中第一維默認是30(3*2*5)。因此，最后get方法返回的是一個Tensor的數據和一個int的標簽。

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(record.path, p)images.extend(seg_imgs)if p < record.num_frames:p += 1process_data = self.transform(images)return process_data, record.labeldef _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]

驗證函數validate基本上和訓練函數train類似，主要有幾個不同點。先是model.eval()將模型設置為evaluate mode，其次沒有optimizer.zero_grad()、loss.backward()、optimizer.step()等損失回傳或梯度更新操作。

def validate(val_loader, model, criterion, iter, logger=None):batch_time = AverageMeter()losses = AverageMeter()top1 = AverageMeter()top5 = AverageMeter()# switch to evaluate modemodel.eval()end = time.time()for i, (input, target) in enumerate(val_loader):target = target.cuda(async=True)input_var = torch.autograd.Variable(input, volatile=True)target_var = torch.autograd.Variable(target, volatile=True)# compute outputoutput = model(input_var)loss = criterion(output, target_var)# measure accuracy and record lossprec1, prec5 = accuracy(output.data, target, topk=(1,5))losses.update(loss.data[0], input.size(0))top1.update(prec1[0], input.size(0))top5.update(prec5[0], input.size(0))# measure elapsed timebatch_time.update(time.time() - end)end = time.time()if i % args.print_freq == 0:print(('Test: [{0}/{1}]\t''Time {batch_time.val:.3f} ({batch_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(i, len(val_loader), batch_time=batch_time, loss=losses,top1=top1, top5=top5)))print(('Testing Results: Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Loss {loss.avg:.5f}'.format(top1=top1, top5=top5, loss=losses)))return top1.avg

準確率計算函數。輸入output是模型預測的結果，尺寸為batch size*num class；target是真實標簽，長度為batch size。這二者都是Tensor類型，具體而言前者是Float Tensor，后者是Long Tensor。batch_size = target.size(0)是讀取batch size值。 _, pred = output.topk(maxk, 1, True, True)這里調用了PyTorch中Tensor的topk方法，第一個輸入maxk表示你要計算的是top maxk的結果；第二個輸入1表示dim，即按行計算（dim=1）；第三個輸入True完整的是largest=True，表示返回的是top maxk個最大值；第四個輸入True完整的是sorted=True，表示返回排序的結果，主要是因為后面要基于這個top maxk的結果計算top 1。target.view(1, -1).expand_as(pred)先將target的尺寸規范到1*batch size，然后將維度擴充為pred相同的維度，也就是maxk*batch size，比如5*batch size，然后調用eq方法計算兩個Tensor矩陣相同元素情況，得到的correct是同等維度的ByteTensor矩陣，1值表示相等，0值表示不相等。correct_k = correct[:k].view(-1).float().sum(0)通過k值來決定是計算top k的準確率，sum(0)表示按照dim 0維度計算和，最后都添加到res列表中并返回。

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].view(-1).float().sum(0)res.append(correct_k.mul_(100.0 / batch_size))return res

總結

以上是生活随笔為你收集整理的TSN算法的PyTorch代码解读（训练部分）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。