File size: 3,282 Bytes
d643072 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
# Copyright 2024 MIT Han Lab
#
# 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
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
import torch
from torch import nn
from .nn.act import build_act, get_act_name
from .nn.conv import ConvLayer
from .nn.norm import build_norm, get_norm_name
from .utils.model import get_same_padding, val2tuple
class MBConvPreGLU(nn.Module):
def __init__(
self,
in_dim: int,
out_dim: int,
kernel_size=3,
stride=1,
mid_dim=None,
expand=6,
padding: int or None = None,
use_bias=False,
norm=(None, None, "ln2d"),
act=("silu", "silu", None),
):
super().__init__()
use_bias = val2tuple(use_bias, 3)
norm = val2tuple(norm, 3)
act = val2tuple(act, 3)
mid_dim = mid_dim or round(in_dim * expand)
self.inverted_conv = ConvLayer(
in_dim,
mid_dim * 2,
1,
use_bias=use_bias[0],
norm=norm[0],
act=None,
)
self.glu_act = build_act(act[0], inplace=False)
self.depth_conv = ConvLayer(
mid_dim,
mid_dim,
kernel_size,
stride=stride,
groups=mid_dim,
padding=padding,
use_bias=use_bias[1],
norm=norm[1],
act=act[1],
)
self.point_conv = ConvLayer(
mid_dim,
out_dim,
1,
use_bias=use_bias[2],
norm=norm[2],
act=act[2],
)
def forward(self, x: torch.Tensor, HW=None) -> torch.Tensor:
B, N, C = x.shape
if HW is None:
H = W = int(N**0.5)
else:
H, W = HW
x = x.reshape(B, H, W, C).permute(0, 3, 1, 2)
x = self.inverted_conv(x)
x, gate = torch.chunk(x, 2, dim=1)
gate = self.glu_act(gate)
x = x * gate
x = self.depth_conv(x)
x = self.point_conv(x)
x = x.reshape(B, C, N).permute(0, 2, 1)
return x
@property
def module_str(self) -> str:
_str = f"{self.depth_conv.kernel_size}{type(self).__name__}("
_str += f"in={self.inverted_conv.in_dim},mid={self.depth_conv.in_dim},out={self.point_conv.out_dim},s={self.depth_conv.stride}"
_str += (
f",norm={get_norm_name(self.inverted_conv.norm)}"
f"+{get_norm_name(self.depth_conv.norm)}"
f"+{get_norm_name(self.point_conv.norm)}"
)
_str += (
f",act={get_act_name(self.inverted_conv.act)}"
f"+{get_act_name(self.depth_conv.act)}"
f"+{get_act_name(self.point_conv.act)}"
)
_str += f",glu_act={get_act_name(self.glu_act)})"
return _str
|