| import torch.nn as nn | |
| from engine.BiRefNet.models.modules.utils import build_act_layer, build_norm_layer | |
| class StemLayer(nn.Module): | |
| r"""Stem layer of InternImage | |
| Args: | |
| in_channels (int): number of input channels | |
| out_channels (int): number of output channels | |
| act_layer (str): activation layer | |
| norm_layer (str): normalization layer | |
| """ | |
| def __init__( | |
| self, | |
| in_channels=3 + 1, | |
| inter_channels=48, | |
| out_channels=96, | |
| act_layer="GELU", | |
| norm_layer="BN", | |
| ): | |
| super().__init__() | |
| self.conv1 = nn.Conv2d( | |
| in_channels, inter_channels, kernel_size=3, stride=1, padding=1 | |
| ) | |
| self.norm1 = build_norm_layer( | |
| inter_channels, norm_layer, "channels_first", "channels_first" | |
| ) | |
| self.act = build_act_layer(act_layer) | |
| self.conv2 = nn.Conv2d( | |
| inter_channels, out_channels, kernel_size=3, stride=1, padding=1 | |
| ) | |
| self.norm2 = build_norm_layer( | |
| out_channels, norm_layer, "channels_first", "channels_first" | |
| ) | |
| def forward(self, x): | |
| x = self.conv1(x) | |
| x = self.norm1(x) | |
| x = self.act(x) | |
| x = self.conv2(x) | |
| x = self.norm2(x) | |
| return x | |