| import torch |
| import torch.nn as nn |
| from functools import partial |
| import torch.nn.functional as F |
|
|
|
|
| class Mlp(nn.Module): |
| def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.,channels_first=True): |
| super().__init__() |
| out_features = out_features or in_features |
| hidden_features = hidden_features or in_features |
|
|
| Linear = partial(nn.Conv2d, kernel_size=1, padding=0) if channels_first else nn.Linear |
| self.fc1 = Linear(in_features, hidden_features) |
| self.act = act_layer() |
| self.fc2 = Linear(hidden_features, out_features) |
| self.drop = nn.Dropout(drop) |
|
|
| def forward(self, x): |
| x = self.fc1(x) |
| x = self.act(x) |
| x = self.drop(x) |
| x = self.fc2(x) |
| x = self.drop(x) |
| return x |
|
|
| def dsconv_3x3(in_channel, out_channel): |
| return nn.Sequential( |
| nn.Conv2d(in_channel, in_channel, kernel_size=3, stride=1, padding=1, groups=in_channel), |
| nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=1, padding=0, groups=1), |
| nn.BatchNorm2d(out_channel), |
| nn.ReLU() |
| ) |
|
|
| class changedetector(nn.Module): |
| def __init__(self, in_channel): |
| super().__init__() |
| self.mlp1 = Mlp(in_features = in_channel, out_features = in_channel) |
| self.mlp2 = Mlp(in_features = in_channel, out_features=2) |
| self.dwc = dsconv_3x3(in_channel, in_channel) |
|
|
| def forward(self, x): |
| x1 = self.mlp1(x) |
| x_d = self.dwc(x1) |
| x_out = self.mlp2(x1 + x_d) |
| x_out = F.interpolate( |
| x_out, |
| scale_factor=(4,4), |
| mode="bilinear", |
| align_corners=False, |
| ) |
| return x_out |
|
|
