VisionLanguageGroup/MicroscopyMatching

1from torch import nn2import torch3 4 5class UpsamplingLayer(nn.Module):6 7    def __init__(self, in_channels, out_channels, leaky=True):8 9        super(UpsamplingLayer, self).__init__()10 11        self.layer = nn.Sequential(12            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),13            nn.LeakyReLU() if leaky else nn.ReLU(),14            nn.UpsamplingBilinear2d(scale_factor=2)15        )16 17    def forward(self, x):18        return self.layer(x)19 20 21class DensityMapRegressor(nn.Module):22 23    def __init__(self, in_channels, reduction):24 25        super(DensityMapRegressor, self).__init__()26 27        if reduction == 8:28            self.regressor = nn.Sequential(29                UpsamplingLayer(in_channels, 128),30                UpsamplingLayer(128, 64),31                UpsamplingLayer(64, 32),32                nn.Conv2d(32, 1, kernel_size=1),33                nn.LeakyReLU()34            )35        elif reduction == 16:36            self.regressor = nn.Sequential(37                UpsamplingLayer(in_channels, 128),38                UpsamplingLayer(128, 64),39                UpsamplingLayer(64, 32),40                UpsamplingLayer(32, 16),41                nn.Conv2d(16, 1, kernel_size=1),42                nn.LeakyReLU()43            )44 45        self.reset_parameters()46 47    def forward(self, x):48        return self.regressor(x)49 50    def reset_parameters(self):51        for module in self.modules():52            if isinstance(module, nn.Conv2d):53                nn.init.normal_(module.weight, std=0.01)54                if module.bias is not None:55                    nn.init.constant_(module.bias, 0)56 57 58