Spaces:
Build error
Build error
from collections import OrderedDict | |
import torch | |
import torch.nn as nn | |
from .bn import ABN | |
class DenseModule(nn.Module): | |
def __init__(self, in_channels, growth, layers, bottleneck_factor=4, norm_act=ABN, dilation=1): | |
super(DenseModule, self).__init__() | |
self.in_channels = in_channels | |
self.growth = growth | |
self.layers = layers | |
self.convs1 = nn.ModuleList() | |
self.convs3 = nn.ModuleList() | |
for i in range(self.layers): | |
self.convs1.append(nn.Sequential(OrderedDict([ | |
("bn", norm_act(in_channels)), | |
("conv", nn.Conv2d(in_channels, self.growth * bottleneck_factor, 1, bias=False)) | |
]))) | |
self.convs3.append(nn.Sequential(OrderedDict([ | |
("bn", norm_act(self.growth * bottleneck_factor)), | |
("conv", nn.Conv2d(self.growth * bottleneck_factor, self.growth, 3, padding=dilation, bias=False, | |
dilation=dilation)) | |
]))) | |
in_channels += self.growth | |
def out_channels(self): | |
return self.in_channels + self.growth * self.layers | |
def forward(self, x): | |
inputs = [x] | |
for i in range(self.layers): | |
x = torch.cat(inputs, dim=1) | |
x = self.convs1[i](x) | |
x = self.convs3[i](x) | |
inputs += [x] | |
return torch.cat(inputs, dim=1) | |