# Copyright (c) OpenMMLab. All rights reserved. import pytest import torch from mmaction.models import build_recognizer from mmaction.utils.gradcam_utils import GradCAM from .base import generate_gradcam_inputs, get_recognizer_cfg def _get_target_shapes(input_shape, num_classes=400, model_type='2D'): if model_type not in ['2D', '3D']: raise ValueError(f'Data type {model_type} is not available') preds_target_shape = (input_shape[0], num_classes) if model_type == '3D': # input shape (batch_size, num_crops*num_clips, C, clip_len, H, W) # target shape (batch_size*num_crops*num_clips, clip_len, H, W, C) blended_imgs_target_shape = (input_shape[0] * input_shape[1], input_shape[3], input_shape[4], input_shape[5], input_shape[2]) else: # input shape (batch_size, num_segments, C, H, W) # target shape (batch_size, num_segments, H, W, C) blended_imgs_target_shape = (input_shape[0], input_shape[1], input_shape[3], input_shape[4], input_shape[2]) return blended_imgs_target_shape, preds_target_shape def _do_test_2D_models(recognizer, target_layer_name, input_shape, num_classes=400, device='cpu'): demo_inputs = generate_gradcam_inputs(input_shape) demo_inputs['imgs'] = demo_inputs['imgs'].to(device) demo_inputs['label'] = demo_inputs['label'].to(device) recognizer = recognizer.to(device) gradcam = GradCAM(recognizer, target_layer_name) blended_imgs_target_shape, preds_target_shape = _get_target_shapes( input_shape, num_classes=num_classes, model_type='2D') blended_imgs, preds = gradcam(demo_inputs) assert blended_imgs.size() == blended_imgs_target_shape assert preds.size() == preds_target_shape blended_imgs, preds = gradcam(demo_inputs, True) assert blended_imgs.size() == blended_imgs_target_shape assert preds.size() == preds_target_shape def _do_test_3D_models(recognizer, target_layer_name, input_shape, num_classes=400): blended_imgs_target_shape, preds_target_shape = _get_target_shapes( input_shape, num_classes=num_classes, model_type='3D') demo_inputs = generate_gradcam_inputs(input_shape, '3D') # parrots 3dconv is only implemented on gpu if torch.__version__ == 'parrots': if torch.cuda.is_available(): recognizer = recognizer.cuda() demo_inputs['imgs'] = demo_inputs['imgs'].cuda() demo_inputs['label'] = demo_inputs['label'].cuda() gradcam = GradCAM(recognizer, target_layer_name) blended_imgs, preds = gradcam(demo_inputs) assert blended_imgs.size() == blended_imgs_target_shape assert preds.size() == preds_target_shape blended_imgs, preds = gradcam(demo_inputs, True) assert blended_imgs.size() == blended_imgs_target_shape assert preds.size() == preds_target_shape else: gradcam = GradCAM(recognizer, target_layer_name) blended_imgs, preds = gradcam(demo_inputs) assert blended_imgs.size() == blended_imgs_target_shape assert preds.size() == preds_target_shape blended_imgs, preds = gradcam(demo_inputs, True) assert blended_imgs.size() == blended_imgs_target_shape assert preds.size() == preds_target_shape def test_tsn(): config = get_recognizer_cfg('tsn/tsn_r50_1x1x3_100e_kinetics400_rgb.py') config.model['backbone']['pretrained'] = None recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 25, 3, 32, 32) target_layer_name = 'backbone/layer4/1/relu' _do_test_2D_models(recognizer, target_layer_name, input_shape) def test_i3d(): config = get_recognizer_cfg('i3d/i3d_r50_32x2x1_100e_kinetics400_rgb.py') config.model['backbone']['pretrained2d'] = False config.model['backbone']['pretrained'] = None recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = [1, 1, 3, 32, 32, 32] target_layer_name = 'backbone/layer4/1/relu' _do_test_3D_models(recognizer, target_layer_name, input_shape) def test_r2plus1d(): config = get_recognizer_cfg( 'r2plus1d/r2plus1d_r34_8x8x1_180e_kinetics400_rgb.py') config.model['backbone']['pretrained2d'] = False config.model['backbone']['pretrained'] = None config.model['backbone']['norm_cfg'] = dict(type='BN3d') recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 3, 3, 8, 32, 32) target_layer_name = 'backbone/layer4/1/relu' _do_test_3D_models(recognizer, target_layer_name, input_shape) def test_slowfast(): config = get_recognizer_cfg( 'slowfast/slowfast_r50_4x16x1_256e_kinetics400_rgb.py') recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 1, 3, 32, 32, 32) target_layer_name = 'backbone/slow_path/layer4/1/relu' _do_test_3D_models(recognizer, target_layer_name, input_shape) def test_tsm(): config = get_recognizer_cfg('tsm/tsm_r50_1x1x8_50e_kinetics400_rgb.py') config.model['backbone']['pretrained'] = None target_layer_name = 'backbone/layer4/1/relu' # base config recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 8, 3, 32, 32) _do_test_2D_models(recognizer, target_layer_name, input_shape) # test twice sample + 3 crops, 2*3*8=48 config.model.test_cfg = dict(average_clips='prob') recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 48, 3, 32, 32) _do_test_2D_models(recognizer, target_layer_name, input_shape) def test_csn(): config = get_recognizer_cfg( 'csn/ircsn_ig65m_pretrained_r152_32x2x1_58e_kinetics400_rgb.py') config.model['backbone']['pretrained2d'] = False config.model['backbone']['pretrained'] = None recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 1, 3, 32, 32, 32) target_layer_name = 'backbone/layer4/1/relu' _do_test_3D_models(recognizer, target_layer_name, input_shape) def test_tpn(): target_layer_name = 'backbone/layer4/1/relu' config = get_recognizer_cfg('tpn/tpn_tsm_r50_1x1x8_150e_sthv1_rgb.py') config.model['backbone']['pretrained'] = None recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 8, 3, 32, 32) _do_test_2D_models(recognizer, target_layer_name, input_shape, 174) config = get_recognizer_cfg( 'tpn/tpn_slowonly_r50_8x8x1_150e_kinetics_rgb.py') config.model['backbone']['pretrained'] = None recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 3, 3, 8, 32, 32) _do_test_3D_models(recognizer, target_layer_name, input_shape) def test_c3d(): config = get_recognizer_cfg('c3d/c3d_sports1m_16x1x1_45e_ucf101_rgb.py') config.model['backbone']['pretrained'] = None recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 1, 3, 16, 112, 112) target_layer_name = 'backbone/conv5a/activate' _do_test_3D_models(recognizer, target_layer_name, input_shape, 101) @pytest.mark.skipif( not torch.cuda.is_available(), reason='requires CUDA support') def test_tin(): config = get_recognizer_cfg( 'tin/tin_tsm_finetune_r50_1x1x8_50e_kinetics400_rgb.py') config.model['backbone']['pretrained'] = None target_layer_name = 'backbone/layer4/1/relu' recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 8, 3, 64, 64) _do_test_2D_models( recognizer, target_layer_name, input_shape, device='cuda:0') def test_x3d(): config = get_recognizer_cfg('x3d/x3d_s_13x6x1_facebook_kinetics400_rgb.py') config.model['backbone']['pretrained'] = None recognizer = build_recognizer(config.model) recognizer.cfg = config input_shape = (1, 1, 3, 13, 32, 32) target_layer_name = 'backbone/layer4/1/relu' _do_test_3D_models(recognizer, target_layer_name, input_shape)