def __init__(self, channels=32, refine=False):
super(SGABlock, self).__init__()
self.refine = refine
if self.refine:
self.bn_relu = nn.Sequential(BatchNorm3d(channels),
nn.ReLU(inplace=True))
self.conv_refine = BasicConv(channels,
channels,
is_3d=True,
kernel_size=3,
padding=1,
relu=False)
# self.conv_refine1 = BasicConv(8, 8, is_3d=True, kernel_size=1, padding=1)
else:
self.bn = BatchNorm3d(channels)
self.SGA = SGA()
self.relu = nn.ReLU(inplace=True)
def __init__(self,
in_channels,
out_channels,
deconv=False,
is_3d=False,
bn=True,
l2=True,
relu=True,
**kwargs):
super(BasicConv, self).__init__()
# print(in_channels, out_channels, deconv, is_3d, bn, relu, kwargs)
self.relu = relu
self.use_bn = bn
self.l2 = l2
if is_3d:
if deconv:
self.conv = nn.ConvTranspose3d(in_channels,
out_channels,
bias=False,
**kwargs)
else:
self.conv = nn.Conv3d(in_channels,
out_channels,
bias=False,
**kwargs)
self.bn = BatchNorm3d(out_channels)
else:
if deconv:
self.conv = nn.ConvTranspose2d(in_channels,
out_channels,
bias=False,
**kwargs)
else:
self.conv = nn.Conv2d(in_channels,
out_channels,
bias=False,
**kwargs)
self.bn = DomainNorm(channel=out_channels, l2=self.l2)
