def __init__(self, nin, nout, nG=64):
super().__init__()
self.name = "UNet"
self.conv0 = nn.Sequential(convBatch(nin, nG),
convBatch(nG, nG))
self.conv1 = nn.Sequential(convBatch(nG * 1, nG * 2, stride=2),
convBatch(nG * 2, nG * 2))
self.conv2 = nn.Sequential(convBatch(nG * 2, nG * 4, stride=2),
convBatch(nG * 4, nG * 4))
self.bridge = nn.Sequential(convBatch(nG * 4, nG * 8, stride=2),
residualConv(nG * 8, nG * 8),
convBatch(nG * 8, nG * 8))
self.deconv1 = upSampleConv(nG * 8, nG * 8)
self.conv5 = nn.Sequential(convBatch(nG * 12, nG * 4),
convBatch(nG * 4, nG * 4))
self.deconv2 = upSampleConv(nG * 4, nG * 4)
self.conv6 = nn.Sequential(convBatch(nG * 6, nG * 2),
convBatch(nG * 2, nG * 2))
self.deconv3 = upSampleConv(nG * 2, nG * 2)
self.conv7 = nn.Sequential(convBatch(nG * 3, nG * 1),
convBatch(nG * 1, nG * 1))
self.final = nn.Conv2d(nG, nout, kernel_size=1)
def __init__(self, nin, nout, nG=32): # m changed the default from 64 to 32
super().__init__()
self.conv0 = nn.Sequential(convBatch(nin, nG),
convBatch(nG, nG))
self.conv1 = nn.Sequential(convBatch(nG * 1, nG * 2, stride=2),
convBatch(nG * 2, nG * 2))
self.conv2 = nn.Sequential(convBatch(nG * 2, nG * 4, stride=2),
convBatch(nG * 4, nG * 4))
self.bridge = nn.Sequential(convBatch(nG * 4, nG * 8, stride=2),
residualConv(nG * 8, nG * 8),
convBatch(nG * 8, nG * 8))
self.deconv1 = upSampleConv(nG * 8, nG * 8)
self.conv5 = nn.Sequential(convBatch(nG * 12, nG * 4),
convBatch(nG * 4, nG * 4))
self.deconv2 = upSampleConv(nG * 4, nG * 4)
self.conv6 = nn.Sequential(convBatch(nG * 6, nG * 2),
convBatch(nG * 2, nG * 2))
self.deconv3 = upSampleConv(nG * 2, nG * 2)
self.conv7 = nn.Sequential(convBatch(nG * 3, nG * 1),
convBatch(nG * 1, nG * 1))
self.final = nn.Conv2d(nG, nout, kernel_size=1)