def __init__(
self,
num_classes=4,
num_channels=3,
dimensions=(128, 128), # pylint: disable=unused-argument
bias=True,
**kwargs
):
super().__init__()
self.enc1 = ai8x.FusedConv2dBNReLU(num_channels, 4, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.enc2 = ai8x.FusedMaxPoolConv2dBNReLU(4, 8, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.enc3 = ai8x.FusedMaxPoolConv2dBNReLU(8, 32, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.bneck = ai8x.FusedMaxPoolConv2dBNReLU(32, 64, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.upconv3 = ai8x.ConvTranspose2d(64, 32, 3, stride=2, padding=1)
self.dec3 = ai8x.FusedConv2dBNReLU(64, 32, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.upconv2 = ai8x.ConvTranspose2d(32, 8, 3, stride=2, padding=1)
self.dec2 = ai8x.FusedConv2dBNReLU(16, 8, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.upconv1 = ai8x.ConvTranspose2d(8, 4, 3, stride=2, padding=1)
self.dec1 = ai8x.FusedConv2dBNReLU(8, 16, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.conv = ai8x.FusedConv2dBN(16, num_classes, 1, stride=1, padding=0,
bias=bias, batchnorm='NoAffine', **kwargs)
def __init__(
self,
num_classes=4,
num_channels=48,
dimensions=(88, 88), # pylint: disable=unused-argument
bias=True,
fold_ratio=4,
**kwargs
):
super().__init__()
self.fold_ratio = fold_ratio
self.num_classes = num_classes
self.num_final_channels = num_classes * fold_ratio * fold_ratio
self.prep0 = ai8x.FusedConv2dBNReLU(num_channels, 64, 1, stride=1, padding=0,
bias=bias, batchnorm='NoAffine', **kwargs)
self.prep1 = ai8x.FusedConv2dBNReLU(64, 64, 1, stride=1, padding=0,
bias=bias, batchnorm='NoAffine', **kwargs)
self.prep2 = ai8x.FusedConv2dBNReLU(64, 32, 1, stride=1, padding=0,
bias=bias, batchnorm='NoAffine', **kwargs)
self.enc1 = ai8x.FusedConv2dBNReLU(32, 8, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.enc2 = ai8x.FusedMaxPoolConv2dBNReLU(8, 28, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.enc3 = ai8x.FusedMaxPoolConv2dBNReLU(28, 56, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.bneck = ai8x.FusedMaxPoolConv2dBNReLU(56, 112, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.upconv3 = ai8x.ConvTranspose2d(112, 56, 3, stride=2, padding=1)
self.dec3 = ai8x.FusedConv2dBNReLU(112, 56, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.upconv2 = ai8x.ConvTranspose2d(56, 28, 3, stride=2, padding=1)
self.dec2 = ai8x.FusedConv2dBNReLU(56, 28, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.upconv1 = ai8x.ConvTranspose2d(28, 8, 3, stride=2, padding=1)
self.dec1 = ai8x.FusedConv2dBNReLU(16, 48, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.dec0 = ai8x.FusedConv2dBNReLU(48, 64, 3, stride=1, padding=1,
bias=bias, batchnorm='NoAffine', **kwargs)
self.conv_p1 = ai8x.FusedConv2dBNReLU(64, 64, 1, stride=1, padding=0,
bias=bias, batchnorm='NoAffine', **kwargs)
self.conv_p2 = ai8x.FusedConv2dBNReLU(64, 64, 1, stride=1, padding=0,
bias=bias, batchnorm='NoAffine', **kwargs)
self.conv_p3 = ai8x.FusedConv2dBN(64, 64, 1, stride=1, padding=0,
bias=bias, batchnorm='NoAffine', **kwargs)
self.conv = ai8x.FusedConv2dBN(64, self.num_final_channels, 1, stride=1, padding=0,
bias=bias, batchnorm='NoAffine', **kwargs)