| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class MultiScaleFusion(nn.Module): |
| def __init__(self, in_ch: int, out_ch: int): |
| super().__init__() |
| self.scale1 = nn.Conv2d(in_ch, in_ch, 3, padding="same") |
| self.scale2 = nn.Conv2d(in_ch, in_ch, 3, padding="same", dilation=2) |
| self.scale3 = nn.Conv2d(in_ch, in_ch, 3, padding="same", dilation=4) |
| self.proj = nn.Conv2d(in_ch * 3, out_ch, 1, bias=False) if in_ch * 3 != out_ch else nn.Identity() |
|
|
| def forward(self, x): |
| f1 = self.scale1(x) |
| f2 = self.scale2(x) |
| f3 = self.scale3(x) |
| fused = torch.cat([f1, f2, f3], dim=1) |
| return self.proj(fused) |
|
|
