def __init__(self, in_size, out_size):
super(SegNetDown2, self).__init__()
self.conv1 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)
self.conv2 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)
self.maxpool_with_argmax = nn.MaxPool2d(2, 2, return_indices=True)
def __init__(self, in_size, out_size):
super(SegNetUp3, self).__init__()
self.unpool = nn.MaxUnpool2d(2, 2)
self.conv1 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)
self.conv2 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)
self.conv3 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)