File size: 1,257 Bytes
7863f92 |
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 |
from mmseg.registry import MODELS
from mmengine.model import BaseModule
from torch import nn as nn
from torch.nn import functional as F
from timm.models.layers import trunc_normal_
@MODELS.register_module()
class FAM(BaseModule):
def __init__(self, in_channels, out_channels, output_size,init_cfg=None):
super().__init__(init_cfg)
self.convert = nn.ModuleList()
self.output_size = output_size
if isinstance(out_channels, int):
out_channels = [out_channels] * len(in_channels)
for in_channel, out_channel in zip(in_channels, out_channels):
self.convert.append(
nn.Conv2d(in_channel, out_channel, kernel_size=1),
)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, (nn.Conv2d, nn.Linear)):
trunc_normal_(m.weight, std=.02)
nn.init.constant_(m.bias, 0)
def forward(self, inputs):
outs = []
for index, x in enumerate(inputs):
x = self.convert[index](x)
x = F.interpolate(
x, size=(self.output_size,self.output_size), align_corners=False, mode="bilinear"
)
outs.append(x)
return tuple(outs) |