动作识别基于PyTorch的3D ResNets模型实现的动作识别任务+运行说明（含训练、微调和测试、在UCF-101和HMDB...

from pathlib import Path import json import random import os import numpy as np import torch from torch.nn import CrossEntropyLoss from torch.optim import SGD, lr_scheduler import torch.multiprocessing as mp import torch.distributed as dist from torch.backends import cudnn import torchvision from opts import parse_opts from model import (generate_model, load_pretrained_model, make_data_parallel, get_fine_tuning_parameters) from mean import get_mean_std from spatial_transforms import (Compose, Normalize, Resize, CenterCrop, CornerCrop, MultiScaleCornerCrop, RandomResizedCrop, RandomHorizontalFlip, ToTensor, ScaleValue, ColorJitter, PickFirstChannels) from temporal_transforms import (LoopPadding, TemporalRandomCrop, TemporalCenterCrop, TemporalEvenCrop, SlidingWindow, TemporalSubsampling) from temporal_transforms import Compose as TemporalCompose from dataset import get_training_data, get_validation_data, get_inference_data from utils import Logger, worker_init_fn, get_lr from training import train_epoch from validation import val_epoch import inference def json_serial(obj): if isinstance(obj, Path): return str(obj) def get_opt(): opt = parse_opts() if opt.root_path is not None: opt.video_path = opt.root_path / opt.video_path opt.annotation_path = opt.root_path / opt.annotation_path opt.result_path = opt.root_path / opt.result_path if opt.resume_path is not None: opt.resume_path = opt.root_path / opt.resume_path if opt.pretrain_path is not None: opt.pretrain_path = opt.root_path / opt.pretrain_path if opt.pretrain_path is not None: opt.n_finetune_classes = opt.n_classes opt.n_classes = opt.n_pretrain_classes if opt.output_topk <= 0: opt.output_topk = opt.n_classes if opt.inference_batch_size == 0: opt.inference_batch_size = opt.batch_size opt.arch = '{}-{}'.format(opt.model, opt.model_depth) opt.begin_epoch = 1 opt.mean, opt.std = get_mean_std(opt.value_scale, dataset=opt.mean_dataset) opt.n_input_channels = 3 if opt.input_type == 'flow': opt.n_input_channels = 2 opt.mean = opt.mean[:2] opt.std = opt.std[:2] if opt.distributed: opt.dist_rank = int(os.environ["OMPI_COMM_WORLD_RANK"]) if opt.dist_rank == 0: print(opt) with (opt.result_path / 'opts.json').open('w') as opt_file: json.dump(vars(opt), opt_file, default=json_serial) else: print(opt) with (opt.result_path / 'opts.json').open('w') as opt_file: json.dump(vars(opt), opt_file, default=json_serial) return opt def resume_model(resume_path, arch, model): print('loading checkpoint {} model'.format(resume_path)) checkpoint = torch.load(resume_path, map_location='cpu') assert arch == checkpoint['arch'] if hasattr(model, 'module'): model.module.load_state_dict(checkpoint['state_dict']) else: model.load_state_dict(checkpoint['state_dict']) return model def resume_train_utils(resume_path, begin_epoch, optimizer, scheduler): print('loading checkpoint {} train utils'.format(resume_path)) checkpoint = torch.load(resume_path, map_location='cpu') begin_epoch = checkpoint['epoch'] + 1 if optimizer is not None and 'optimizer' in checkpoint: optimizer.load_state_dict(checkpoint['optimizer']) if scheduler is not None and 'scheduler' in checkpoint: scheduler.load_state_dict(checkpoint['scheduler']) return begin_epoch, optimizer, scheduler def get_normalize_method(mean, std, no_mean_norm, no_std_norm): if no_mean_norm: if no_std_norm: return Normalize([0, 0, 0], [1, 1, 1]) else: return Normalize([0, 0, 0], std) else: if no_std_norm: return Normalize(mean, [1, 1, 1]) else: return Normalize(mean, std) def get_train_utils(opt, model_parameters): assert opt.train_crop in ['random', 'corner', 'center'] spatial_transform = [] if opt.train_crop == 'random': spatial_transform.append( RandomResizedCrop( opt.sample_size, (opt.train_crop_min_scale, 1.0), (opt.train_crop_min_ratio, 1.0 / opt.train_crop_min_ratio))) elif opt.train_crop == 'corner': scales = [1.0] scale_step = 1 / (2**(1 / 4)) for _ in range(1, 5): scales.append(scales[-1] * scale_step) spatial_transform.append(MultiScaleCornerCrop(opt.sample_size, scales)) elif opt.train_crop == 'center': spatial_transform.append(Resize(opt.sample_size)) spatial_transform.append(CenterCrop(opt.sample_size)) normalize = get_normalize_method(opt.mean, opt.std, opt.no_mean_norm, opt.no_std_norm) if not opt.no_hflip: spatial_transform.append(RandomHorizontalFlip()) if opt.colorjitter: spatial_transform.append(ColorJitter()) spatial_transform.append(ToTensor()) if opt.input_type == 'flow': spatial_transform.append(PickFirstChannels(n=2)) spatial_transform.append(ScaleValue(opt.value_scale)) spatial_transform.append(normalize) spatial_transform = Compose(spatial_transform) assert opt.train_t_crop in ['random', 'center'] temporal_transform = [] if opt.sample_t_stride > 1: temporal_transform.append(TemporalSubsampling(opt.sample_t_stride)) if opt.train_t_crop == 'random': temporal_transform.append(TemporalRandomCrop(opt.sample_duration)) elif opt.train_t_crop == 'center': temporal_transform.append(TemporalCenterCrop(opt.sample_duration)) temporal_transform = TemporalCompose(temporal_transform) train_data = get_training_data(opt.video_path, opt.annotation_path, opt.dataset, opt.input_type, opt.file_type, spatial_transform, temporal_transform) if opt.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_data) else: train_sampler = None train_loader = torch.utils.data.DataLoader(train_data, batch_size=opt.batch_size, shuffle=(train_sampler is None), num_workers=opt.n_threads, pin_memory=True, sampler=train_sampler, worker_init_fn=worker_init_fn) if opt.is_master_node: train_logger = Logger(opt.result_path / 'train.log', ['epoch', 'loss', 'acc', 'lr']) train_batch_logger = Logger( opt.result_path / 'train_batch.log', ['epoch', 'batch', 'iter', 'loss', 'acc', 'lr']) else: train_logger = None train_batch_logger = None if opt.nesterov: dampening = 0 else: dampening = opt.dampening optimizer = SGD(model_parameters, lr=opt.learning_rate, momentum=opt.momentum, dampening=dampening, weight_decay=opt.weight_decay, nesterov=opt.nesterov) assert opt.lr_scheduler in ['plateau', 'multistep'] assert not (opt.lr_scheduler == 'plateau' and opt.no_val) if opt.lr_scheduler == 'plateau': scheduler = lr_scheduler.ReduceLROnPlateau( optimizer, 'min', patience=opt.plateau_patience) else: scheduler = lr_scheduler.MultiStepLR(optimizer, opt.multistep_miles