I’m trying to create a model that has a convolutional and deconvolutional pathway, with tied weights. Essentially, the deconv path is the mirror image of the conv path. Similar to this paper Deconvolutional Feature Stacking… 1.
I like being able to use the nn.Module methods and attributes, but I wanted to create the weights directly, and pass them to the F.conv3d and F.conv3d_transpose functions, just because it seemed less straightforward to keep the weights tied using nn.Conv3d.
this was my attempt, but something is definitely wrong.
class Model3d(nn.Module):

    def __init__(self, input_channels=1, n_classes=2, img_size=96, n_feature_maps=(8, 12, 16, 20, 20)):

        super(Model3d, self).__init__()
        self.input_channels = input_channels
        self.n_feat_map = n_feature_maps
        self.n_classes = n_classes
        self.img_size = img_size

        self.weight1 = self.weight_variable([self.n_feat_map[0], input_channels, 3, 3, 3])
        self.bias_c_1 = self.bias_variable(self.n_feat_map[0])
        self.bias_d_5 = self.bias_variable(input_channels)

        self.weight2 = self.weight_variable([self.n_feat_map[1], self.n_feat_map[0], 3, 3, 3])
        self.bias_c_2 = self.bias_variable(self.n_feat_map[1])
        self.bias_d_4 = self.bias_variable(self.n_feat_map[0])

        self.weight3 = self.weight_variable([self.n_feat_map[2], self.n_feat_map[1], 3, 3, 3])
        self.bias_c_3 = self.bias_variable(self.n_feat_map[2])
        self.bias_d_3 = self.bias_variable(self.n_feat_map[1])

        self.weight4 = self.weight_variable([self.n_feat_map[3], self.n_feat_map[2], 3, 3, 3])
        self.bias_c_4 = self.bias_variable(self.n_feat_map[3])
        self.bias_d_2 = self.bias_variable(self.n_feat_map[2])

        self.weight5 = self.weight_variable([self.n_feat_map[4], self.n_feat_map[3], 3, 3, 3])
        self.bias_c_5 = self.bias_variable(self.n_feat_map[4])
        self.bias_d_1 = self.bias_variable(self.n_feat_map[3])

        self.down_classifier_weight = self.weight_variable([self.n_classes, self.n_feat_map[4], 1, 1, 1])

        self.down_classifier_bias = self.bias_variable(self.n_classes)
        self.up_classifier_weight = self.weight_variable([self.n_classes,
                                                     np.sum(self.n_feat_map)+self.input_channels, 1, 1, 1])
        self.up_classifier_bias = self.bias_variable(self.n_classes)

    def forward(self, x):

        x, ind1 = down_block3d(x, self.weight1, self.bias_c_1)
        x, ind2 = down_block3d(x, self.weight2, self.bias_c_2)
        x, ind3 = down_block3d(x, self.weight3, self.bias_c_3)
        x, ind4 = down_block3d(x, self.weight4, self.bias_c_4)
        down5, ind5 = down_block3d(x, self.weight5, self.bias_c_5)

        up1 = up_block3d(down5, self.weight5, self.bias_d_1, ind5)
        up2 = up_block3d(up1, self.weight4, self.bias_d_2, ind4)
        up3 = up_block3d(up2, self.weight3, self.bias_d_3, ind3)
        up4 = up_block3d(up3, self.weight2, self.bias_d_4, ind2)
        up5 = up_block3d(up4, self.weight1, self.bias_d_5, ind1)

        classifier = classifier3d(down5, self.down_classifier_weight, self.down_classifier_bias)
        to_stack = (down5, up1, up2, up3, up4)

        upsampled = [normalize_and_upsample(i, self.img_size) for i in to_stack] + [up5]
        segmap = classifier3d(torch.cat(upsampled, 1), self.up_classifier_weight, self.up_classifier_bias)

        return classifier, segmap

    def weight_variable(self, shape):
        var =  nn.Parameter(torch.zeros(shape))
        xavier_uniform_(var)
        return var

    def bias_variable(self, size):
        var = nn.Parameter(torch.zeros(size))
        return var

def down_block3d(input, weights, bias):
    x = F.conv3d(input, weights, bias=bias, stride=1, padding=1)
    x = F.relu(x)
    out, indices = F.max_pool3d(x,2,stride=2,return_indices=True)
    return out, indices

def up_block3d(input, weights, bias, indices):
    x = F.max_unpool3d(input, indices, 2, stride=2)
    pad = (-1,-1,-1,-1,-1,-1)
    padded = F.pad(x, pad)
    x = F.conv_transpose3d(padded, weights, bias=bias)
    return F.relu(x)

def classifier3d(input, weights, bias):
    x = F.conv3d(input, weights, bias)
    return F.relu(x)

def normalize_channel_wise(x):
    x = x - torch.mean(x, 1, keepdim=True)
    return x / torch.std(x, 1, keepdim=True)

def normalize_and_upsample(x, new_size):
    x = normalize_channel_wise(x)
    return F.upsample(x, size=new_size)

First of all, in all the examples I have seen, nn.Module subclasses are callable. Mine is not callable, I have been explicitly calling model.forward(). It also seems that gradients are not being propagated, as when I call loss.backward(), all the gradients go to nan, and the GPU is not being utilized at all during the training loop, even though the memory has been allocated, and I have called .cuda() on all the inputs and the model.
I’m wondering, is there an easy to fix my work? or, is there a way to re-use the conv kernels like this with nn.Conv3d and the equivalent Transposed op?