model_builder.py

Go to the documentation of this file.

from smaction.models.backbones.stgcn import STGCN
from smaction.models.heads.gcn_head import GCNHead
from smaction.models.backbones.resnet3d import ResNet3d
from smaction.models.heads.i3d_head import I3DHead
from smaction.models.heads.i3d_fusion import I3DFusion
from smaction.models.heads.linear_head import LinearHead
from smaction.models.heads.mlp_head import MLPHead
from smaction.models.backbones.resnet2d import ResNet2d
from smaction.models.heads.cat_layer import CatLayer
 
 
__model_builders__ = {
    "STGCN" : STGCN,
    "GCNHead" : GCNHead,
    "ResNet3d" : ResNet3d,
    "I3DHead" : I3DHead,
    "I3DFusion" : I3DFusion,
    "LinearHead" : LinearHead,
    "MLPHead" : MLPHead,
    "ResNet2d" : ResNet2d,
    "CatLayer" : CatLayer,
}
 
def build_action_model(args):
    """
    action/models 의 모듈을 빌드하는 함수
 
    args : 
        type -> 모듈의 타입  __model_builders__의 키값이 들어가야함
        type을 제외한 나머지 파라미터는 모듈의 파라미터값이 있어야함
    
    returns:
        파이토이 모델(모듈)
    """
    assert args.type in __model_builders__, \
            f'not found action model type : {args.type}'
    
    runner_builder = __model_builders__[args.type]
    runner_args = args.copy()
    runner_args.pop('type')
    return runner_builder(**runner_args)