這篇部落格來讀一讀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 dataset import TSNDataSet
from models import TSN
from transforms import *
from opts import parser
best_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_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 [email protected] 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'
'[email protected] {top1.val:.3f} ({top1.avg:.3f})\t'
'[email protected] {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的長度乘以(52)。process_data = self.transform(images)将list類型的images封裝成了Tensor，在訓練的時候：對于RGB輸入，這個Tensor的尺寸是(3self.num_segments,224,224)，其中3表示3通道彩色；對于Flow輸入，這個Tensor的尺寸是(self.num_segments2self.new_length,224,224)，其中第一維預設是30(325)。是以，最後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'
'[email protected] {top1.val:.3f} ({top1.avg:.3f})\t'
'[email protected] {top5.val:.3f} ({top5.avg:.3f})'.format(
i, len(val_loader), batch_time=batch_time, loss=losses,
top1=top1, top5=top5)))
print(('Testing Results: [email protected] {top1.avg:.3f} [email protected] {top5.avg:.3f} Loss {loss.avg:.5f}'
.format(top1=top1, top5=top5, loss=losses)))
return top1.avg
準确率計算函數。輸入output是模型預測的結果，尺寸為batch sizenum 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的尺寸規範到1batch size，然後将次元擴充為pred相同的次元，也就是maxkbatch size，比如5batch 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 [email protected] 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
作者：AI之路
來源：CSDN
原文：https://blog.csdn.net/u014380165/article/details/79058147
版權聲明：本文為部落客原創文章，轉載請附上博文連結！
![PyTorch中Tensor的學習筆記1 緻謝2 如何建立一個Tensor2.2 如何建立作為參數的Tensor3 PyTorch中Tensor的索引[圖]](data:image/gif;base64,R0lGODlhAQABAIAAAP///wAAACwAAAAAAQABAAACAkQBADs=)