TSN演算法的PyTorch程式碼解讀（訓練部分）
這篇部落格來讀一讀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_norm
from main函式主要包含匯入模型、資料準備、訓練三個部分，接下來將按順序介紹。parser是在opts.py中定義的關於讀取命令列引數的物件，然後通過from opts import parser匯入的。model = TSN(num_class, args.num_segments, args.modality,...，partial_bn=not args.no_partialbn)
def main():
global args, best_prec1
args = parser.parse_args()
if args.dataset == 'ucf101':
num_class = 101
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)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 5
else:
self.new_length = new_length
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)_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_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)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_channel
setattr(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]] = channel
self.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_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前面提到如果輸入不是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 manner
modules = 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 kernels
params = [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_kernels
if len(params) == 2:
new_conv.bias.data = params[1].data # add bias if neccessary
layer_name = list(container.state_dict().keys())[0][:-7] # remove .weight suffix to get the layer name
# replace the first convlution layer
setattr(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_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)))
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 code
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("", 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_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()
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 optimizer
if 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)
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, 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 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(input_var)
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1, 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 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'])))前面提到在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 = 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前面提到在執行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_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在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 += 1
process_data = self.transform(images)
return process_data, record.label
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]驗證函式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 mode
model.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 output
output = model(input_var)
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1, 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 time
batch_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