def __init__(self, backbone, cls_head, neck=None, train_cfg=None, test_cfg=None): super().__init__() # The backbones in mmcls can be used by TSN if backbone['type'].startswith('mmcls.'): try: import mmcls.models.builder as mmcls_builder except (ImportError, ModuleNotFoundError): raise ImportError('Please install mmcls to use this backbone.') backbone['type'] = backbone['type'][6:] self.backbone = mmcls_builder.build_backbone(backbone) else: self.backbone = builder.build_backbone(backbone) if neck is not None: self.neck = builder.build_neck(neck) self.cls_head = builder.build_head(cls_head) self.train_cfg = train_cfg self.test_cfg = test_cfg # aux_info is the list of tensor names beyond 'imgs' and 'label' which # will be used in train_step and val_step, data_batch should contain # these tensors self.aux_info = [] if train_cfg is not None and 'aux_info' in train_cfg: self.aux_info = train_cfg['aux_info'] # max_testing_views should be int self.max_testing_views = None if test_cfg is not None and 'max_testing_views' in test_cfg: self.max_testing_views = test_cfg['max_testing_views'] assert isinstance(self.max_testing_views, int) self.init_weights() self.fp16_enabled = False
def __init__(self, backbone, cls_head, neck=None, train_cfg=None, test_cfg=None): super().__init__() # record the source of the backbone self.backbone_from = 'mmaction2' if backbone['type'].startswith('mmcls.'): try: import mmcls.models.builder as mmcls_builder except (ImportError, ModuleNotFoundError): raise ImportError('Please install mmcls to use this backbone.') backbone['type'] = backbone['type'][6:] self.backbone = mmcls_builder.build_backbone(backbone) self.backbone_from = 'mmcls' elif backbone['type'].startswith('torchvision.'): try: import torchvision.models except (ImportError, ModuleNotFoundError): raise ImportError('Please install torchvision to use this ' 'backbone.') backbone_type = backbone.pop('type')[12:] self.backbone = torchvision.models.__dict__[backbone_type]( **backbone) # disable the classifier self.backbone.classifier = nn.Identity() self.backbone.fc = nn.Identity() self.backbone_from = 'torchvision' else: self.backbone = builder.build_backbone(backbone) if neck is not None: self.neck = builder.build_neck(neck) self.cls_head = builder.build_head(cls_head) self.train_cfg = train_cfg self.test_cfg = test_cfg # aux_info is the list of tensor names beyond 'imgs' and 'label' which # will be used in train_step and val_step, data_batch should contain # these tensors self.aux_info = [] if train_cfg is not None and 'aux_info' in train_cfg: self.aux_info = train_cfg['aux_info'] # max_testing_views should be int self.max_testing_views = None if test_cfg is not None and 'max_testing_views' in test_cfg: self.max_testing_views = test_cfg['max_testing_views'] assert isinstance(self.max_testing_views, int) # mini-batch blending, e.g. mixup, cutmix, etc. self.blending = None if train_cfg is not None and 'blending' in train_cfg: from mmcv.utils import build_from_cfg from mmaction.datasets.builder import BLENDINGS self.blending = build_from_cfg(train_cfg['blending'], BLENDINGS) self.init_weights() self.fp16_enabled = False