a small suggestion to SegResNet

[function2-llx]

Hi there,

If I'm understanding correctly, SegResNet is a ResNet encoder + U-Net decoder, so how about making SegResNet return the result of the encoder in forward? This would be useful if running classification and segmentation at the same time. So does SegResNetVAE.

Hope to hear anyone's opinions on this!

[function2-llx]

Okay, I would say that it's enough for SegResNet to call the encode and decode methods successively. However, there seems not to have a convenient interface for SegResNetVAE.

[function2-llx]

BTW, seems that the implementation of the forward method of SegResNetVAE could reuse the method from SegResNet.
Currently (according to code in repo) it is:

def forward(self, x):
    net_input = x
    x = self.convInit(x)
    if self.dropout_prob is not None:
        x = self.dropout(x)

    down_x = []
    for down in self.down_layers:
        x = down(x)
        down_x.append(x)

    down_x.reverse()

    vae_input = x

    for i, (up, upl) in enumerate(zip(self.up_samples, self.up_layers)):
        x = up(x) + down_x[i + 1]
        x = upl(x)

    if self.use_conv_final:
        x = self.conv_final(x)

    if self.training:
        vae_loss = self._get_vae_loss(net_input, vae_input)
        return x, vae_loss

    return x, None

which could be rewritten with something like:

def forward(self, x):
    net_input = x

    x, down_x = self.encode(x)
    down_x.reverse()

    vae_input = x

    x = self.decode(x, down_x)

    if self.training:
        vae_loss = self._get_vae_loss(net_input, vae_input)
        return x, vae_loss

    return x, None

[wyli]

thanks, this is addressed by 8a454bd