Source code for towhee.models.movinet.movinet

# Copyright 2022 Zilliz. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
# Code inspired by
# Original code from
# Modified by Zilliz.

from collections import OrderedDict
import torch
import torch.nn.functional as F
from torch import nn, Tensor
from towhee.models.movinet.movinet_block import ConvBlock3D, BasicBneck
from towhee.models.utils.causal_module import CausalModule
from towhee.models.layers.activations import Swish
from towhee.models.layers.temporal_cg_avgpool3d import TemporalCGAvgPool3D
from towhee.models.movinet.config import _C

[docs]class MoViNet(nn.Module): """ Args: causal(`string`): Causal mode. pretrained(`bool`): Pretrained models. If pretrained is True, num_classes is set to 600, conv_type is set to "3d" if causal is False, "2plus1d" if causal is True tf_like is set to True. num_classes(`int`): Number of classes for classifcation. conv_type(`string`): Type of convolution either 3d or 2plus1d tf_like(`bool`): Tf_like behaviour, basically same padding for convolutions. """
[docs] def __init__(self, cfg: "CfgNode", causal: bool = True, pretrained: bool = False, num_classes: int = 600, conv_type: str = "3d", tf_like: bool = False ) -> None: super().__init__() if pretrained: tf_like = True num_classes = 600 conv_type = "2plus1d" if causal else "3d" blocks_dic = OrderedDict() norm_layer = nn.BatchNorm3d if conv_type == "3d" else nn.BatchNorm2d activation_layer = Swish if conv_type == "3d" else nn.Hardswish # conv1 self.conv1 = ConvBlock3D( in_planes=cfg.conv1.input_channels, out_planes=cfg.conv1.out_channels, kernel_size=cfg.conv1.kernel_size, stride=cfg.conv1.stride, padding=cfg.conv1.padding, causal=causal, conv_type=conv_type, tf_like=tf_like, norm_layer=norm_layer, activation_layer=activation_layer ) # blocks for i, block in enumerate(cfg.blocks): for j, basicblock in enumerate(block): blocks_dic[f"b{i}_l{j}"] = BasicBneck(basicblock, causal=causal, conv_type=conv_type, tf_like=tf_like, norm_layer=norm_layer, activation_layer=activation_layer ) self.blocks = nn.Sequential(blocks_dic) # conv7 self.conv7 = ConvBlock3D( in_planes=cfg.conv7.input_channels, out_planes=cfg.conv7.out_channels, kernel_size=cfg.conv7.kernel_size, stride=cfg.conv7.stride, padding=cfg.conv7.padding, causal=causal, conv_type=conv_type, tf_like=tf_like, norm_layer=norm_layer, activation_layer=activation_layer ) # pool self.classifier = nn.Sequential( # dense9 ConvBlock3D(cfg.conv7.out_channels, cfg.dense9.hidden_dim, kernel_size=(1, 1, 1), tf_like=tf_like, causal=causal, conv_type=conv_type, bias=True), Swish(), nn.Dropout(p=0.2, inplace=True), # dense10d ConvBlock3D(cfg.dense9.hidden_dim, num_classes, kernel_size=(1, 1, 1), tf_like=tf_like, causal=causal, conv_type=conv_type, bias=True), ) if causal: self.cgap = TemporalCGAvgPool3D() if pretrained: if causal: if not in ["A0", "A1", "A2"]: raise ValueError("Only A0,A1,A2 streaming" + "networks are available pretrained") state_dict = (torch.hub .load_state_dict_from_url(cfg.stream_weights)) else: state_dict = torch.hub.load_state_dict_from_url(cfg.weights) self.load_state_dict(state_dict) else: self.apply(self._weight_init) self.causal = causal
def avg(self, x: Tensor) -> Tensor: if self.causal: avg = F.adaptive_avg_pool3d(x, (x.shape[2], 1, 1)) avg = self.cgap(avg)[:, :, -1:] else: avg = F.adaptive_avg_pool3d(x, 1) return avg @staticmethod def _weight_init(m): if isinstance(m, nn.Conv3d): nn.init.kaiming_normal_(m.weight, mode="fan_out") if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, (nn.BatchNorm3d, nn.BatchNorm2d, nn.GroupNorm)): nn.init.ones_(m.weight) nn.init.zeros_(m.bias) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) nn.init.zeros_(m.bias) def forward_features(self, x: Tensor) -> Tensor: x = self.conv1(x) x = self.blocks(x) x = self.conv7(x) x = self.avg(x) return x def head(self, x: Tensor) -> Tensor: x = self.classifier(x) x = x.flatten(1) return x def _forward_impl(self, x: Tensor) -> Tensor: x = self.forward_features(x) x = self.head(x) return x
[docs] def forward(self, x: Tensor) -> Tensor: return self._forward_impl(x)
@staticmethod def _clean_activation_buffers(m): if issubclass(type(m), CausalModule): m.reset_activation() def clean_activation_buffers(self) -> None: self.apply(self._clean_activation_buffers)
[docs]def create_model( model_name: str = "movineta0", pretrained: bool = False, causal: bool = False, device: str = None, ): if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" if model_name == "movineta0": model_config = _C.MODEL.MoViNetA0 elif model_name == "movineta1": model_config = _C.MODEL.MoViNetA1 elif model_name == "movineta2": model_config = _C.MODEL.MoViNetA2 elif model_name == "movineta3": model_config = _C.MODEL.MoViNetA3 elif model_name == "movineta4": model_config = _C.MODEL.MoViNetA4 elif model_name == "movineta5": model_config = _C.MODEL.MoViNetA5 else: raise AttributeError(f"Invalid model_name {model_name}.") model = MoViNet( cfg = model_config, causal = causal, pretrained = pretrained, num_classes = 600, conv_type = "3d", tf_like = False ) return model