in pytorchvideo/accelerator/deployment/mobile_cpu/transmuter/transmuter_mobile_cpu.py [0:0]
def transmute_Conv3d5x1x1BnAct(input_module: nn.Module):
"""
Given an input_module, transmutes it into a equivalent Conv3d5x1x1BnAct.
Returns None if no equivalent Conv3d5x1x1BnAct is found, else returns
an instance of equivalent Conv3d5x1x1BnAct.
Args:
input_module (nn.Module): input module to find an equivalent Conv3d5x1x1BnAct
"""
if not isinstance(input_module, nn.Conv3d):
return None
if (
input_module.kernel_size == (5, 1, 1)
and input_module.stride == (1, 1, 1)
and input_module.padding == (2, 0, 0)
and input_module.dilation == (1, 1, 1)
and input_module.padding_mode == "zeros"
):
module = Conv3d5x1x1BnAct(
in_channels=input_module.in_channels,
out_channels=input_module.out_channels,
bias=False if input_module.bias is None else True,
groups=input_module.groups,
activation="identity",
use_bn=False,
)
module.kernel.conv.load_state_dict(input_module.state_dict())
return module
else:
return None