def load_train_config(args): if args.config_file: cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) if args.log_step != -1: cfg.TRAIN.LOG_STEP = args.log_step if args.save_step != -1: cfg.TRAIN.SAVE_STEP = args.save_step if args.eval_step != -1: cfg.TRAIN.EVAL_STEP = args.eval_step if args.resume: cfg.TRAIN.RESUME = True if not args.use_tensorboard: cfg.TRAIN.USE_TENSORBOARD = False if args.gpus != -1: cfg.NUM_GPUS = args.gpus if args.nodes != -1: cfg.NUM_NODES = args.nodes if args.nr != -1: cfg.RANK_ID = args.nr num_gpus = cfg.NUM_GPUS if num_gpus > 1: cfg.OPTIMIZER.LR *= num_gpus cfg.freeze() if not os.path.exists(cfg.OUTPUT_DIR): os.makedirs(cfg.OUTPUT_DIR) return cfg
def test_shufflenet(): cfg.merge_from_file('configs/tsn_sfv2_ucf101_rgb_raw_dense_1x16x4.yaml') model = build_shufflenet_v2(cfg) data = torch.randn((1, 3, 224, 224)) outputs = model(data) print(outputs.shape) assert outputs.shape == (1, 1024, 7, 7)
def test_resnet50(): cfg.merge_from_file('configs/tsn_r50_ucf101_rgb_raw_dense_1x16x4.yaml') model = build_resnet50(cfg) data = torch.randn((1, 3, 224, 224)) outputs = model(data) print(outputs.shape) assert outputs.shape == (1, 2048, 7, 7)
def test_jester_rgbdiff(): cfg.merge_from_file('configs/tsn_r50_jester_rgbdiff_224x3_seg.yaml') transform = build_transform(cfg, is_train=True) dataset = build_dataset(cfg, transform=transform, is_train=True) image, target = dataset.__getitem__(20) print(image.shape) print(target) assert image.shape == (3, 15, 224, 224)
def test_hmdb51_rgb(): cfg.merge_from_file('configs/tsn_r50_hmdb51_rgb_224x3_seg.yaml') cfg.DATASETS.NUM_CLIPS = 8 transform = build_transform(cfg, is_train=True) dataset = build_dataset(cfg, transform=transform, is_train=True) image, target = dataset.__getitem__(20) print(image.shape) print(target) assert image.shape == (3, 8, 224, 224)
def test_slowonly(): cfg.merge_from_file('configs/slowonly_r3d50_ucf101_rgb_224x4_dense.yaml') cfg.freeze() model = resnet3d_50_slowonly(cfg) print(model) data = torch.randn(1, 3, 4, 224, 224) outputs = model(data) print(outputs.shape) assert outputs.shape == (1, 2048, 4, 7, 7)
def test_i3d(): data = torch.randn(1, 3, 32, 224, 224) cfg.merge_from_file( 'configs/i3d-3x1-nl_r3d50_ucf101_rgb_224x32_dense.yaml') model = resnet3d_50(cfg) outputs = model(data) print(outputs.shape) assert outputs.shape == (1, 2048, 4, 7, 7) cfg.merge_from_file('configs/i3d-3x1_r3d50_ucf101_rgb_224x32_dense.yaml') model = resnet3d_50(cfg) outputs = model(data) print(outputs.shape) assert outputs.shape == (1, 2048, 4, 7, 7) cfg.merge_from_file( 'configs/i3d-3x3-nl_r3d50_ucf101_rgb_224x32_dense.yaml') model = resnet3d_50(cfg) outputs = model(data) print(outputs.shape) assert outputs.shape == (1, 2048, 4, 7, 7) cfg.merge_from_file('configs/i3d-3x3_r3d50_ucf101_rgb_224x32_dense.yaml') model = resnet3d_50(cfg) outputs = model(data) print(outputs.shape) assert outputs.shape == (1, 2048, 4, 7, 7)
def load_config(args): if args.config_file: cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) if args.gpus > 1: cfg.OPTIMIZER.LR *= args.gpus cfg.OPTIMIZER.WEIGHT_DECAY *= args.gpus cfg.LR_SCHEDULER.COSINE_ANNEALING_LR.MINIMAL_LR *= args.gpus cfg.freeze() if not os.path.exists(cfg.OUTPUT.DIR): os.makedirs(cfg.OUTPUT.DIR) return cfg
def test_slowfast(): cfg.merge_from_file('configs/slowfast_r3d50_ucf101_rgb_224x32_dense.yaml') cfg.freeze() model = resnet3d_50_slowfast(cfg) print(model) data = torch.randn(1, 3, 32, 224, 224) outputs = model(data) print(len(outputs)) print(outputs[0].shape) print(outputs[1].shape) assert len(outputs) == 2 assert outputs[0].shape == (1, 2048, 4, 7, 7) assert outputs[1].shape == (1, 256, 32, 7, 7)
def load_test_config(args): if not os.path.isfile(args.config_file) or not os.path.isfile(args.pretrained): raise ValueError('需要输入配置文件和预训练模型路径') cfg.merge_from_file(args.config_file) cfg.MODEL.PRETRAINED = args.pretrained cfg.OUTPUT_DIR = args.output if args.gpus != -1: cfg.NUM_GPUS = args.gpus if args.nodes != -1: cfg.NODES = args.nodes if args.nr != -1: cfg.RANK = args.nr cfg.freeze() if not os.path.exists(cfg.OUTPUT_DIR): os.makedirs(cfg.OUTPUT_DIR) return cfg
def main(data_shape, config_file, mobile_name): cfg.merge_from_file(config_file) device = get_device(local_rank=get_local_rank()) model = build_recognizer(cfg, device) data = torch.randn(data_shape).to(device=device, non_blocking=True) GFlops, params_size = compute_num_flops(model, data) print(f'{mobile_name} ' + '*' * 10) print(f'device: {device}') print(f'GFlops: {GFlops}') print(f'Params Size: {params_size}') total_time = 0.0 num = 100 for i in range(num): data = torch.randn(data_shape).to(device=device, non_blocking=True) start = time.time() model(data) total_time = time.time() - start print(f'one process need {total_time / num}')
def main(): parser = argparse.ArgumentParser(description='TSN Test With PyTorch') parser.add_argument("config_file", default="", metavar="CONFIG_FILE", help="path to config file", type=str) parser.add_argument('pretrained', default="", metavar='PRETRAINED_FILE', help="path to pretrained model", type=str) parser.add_argument('--output', default="./outputs/test", type=str) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() if not os.path.isfile(args.config_file) or not os.path.isfile(args.pretrained): raise ValueError('需要输入配置文件和预训练模型路径') cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.MODEL.PRETRAINED = args.pretrained cfg.OUTPUT.DIR = args.output cfg.freeze() if not os.path.exists(cfg.OUTPUT.DIR): os.makedirs(cfg.OUTPUT.DIR) logger = setup_logger("TSN", save_dir=cfg.OUTPUT.DIR) logger.info(args) logger.info("Environment info:\n" + collect_env_info()) logger.info("Loaded configuration file {}".format(args.config_file)) if args.config_file: with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) test(cfg)
def main(data_shape, config_file, mobile_name): cfg.merge_from_file(config_file) device = get_device(local_rank=get_local_rank()) model = build_recognizer(cfg, device) model.eval() data = torch.randn(data_shape).to(device=device, non_blocking=True) GFlops, params_size = compute_num_flops(model, data) print(f'{mobile_name} ' + '*' * 10) print(f'device: {device}') print(f'GFlops: {GFlops}') print(f'Params Size: {params_size}') data = torch.randn(data_shape) t1 = 0.0 num = 100 begin = time.time() for i in range(num): start = time.time() model(data.to(device=device, non_blocking=True)) t1 += time.time() - start t2 = time.time() - begin print(f'one process need {t2 / num}, model compute need: {t1 / num}')
def main(): parser = argparse.ArgumentParser(description='TSN Training With PyTorch') parser.add_argument("--config_file", default="", metavar="FILE", help="path to config file", type=str) parser.add_argument('--log_step', default=10, type=int, help='Print logs every log_step') parser.add_argument('--save_step', default=2500, type=int, help='Save checkpoint every save_step') parser.add_argument('--stop_save', default=False, action='store_true') parser.add_argument( '--eval_step', default=2500, type=int, help='Evaluate dataset every eval_step, disabled when eval_step < 0') parser.add_argument('--stop_eval', default=False, action='store_true') parser.add_argument('--resume', default=False, action='store_true', help='Resume training') parser.add_argument('--use_tensorboard', default=1, type=int) parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N', help='number of machines (default: 1)') parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus per node') parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes') parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() if args.config_file: cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() if not os.path.exists(cfg.OUTPUT.DIR): os.makedirs(cfg.OUTPUT.DIR) logger = setup_logger("TSN", save_dir=cfg.OUTPUT.DIR) logger.info(args) logger.info("Environment info:\n" + collect_env_info()) logger.info("Loaded configuration file {}".format(args.config_file)) if args.config_file: with open(args.config_file, "r") as cf: config_str = "\n" + cf.read() logger.info(config_str) logger.info("Running with config:\n{}".format(cfg)) args.world_size = args.gpus * args.nodes os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '14028' mp.spawn(train, nprocs=args.gpus, args=(args, cfg))