Is it possible to use Unet for multi-output prediction?

[Minxiangliu]

I have some PET images and have labeled the ROI of the cancer range, however I want to use this data to make other predictions, including probabilistic prediction, linear regression, etc., so there is a way to do like what this blog mentions Multiple outputs and multiple losses architecture?

Thanks in advance!

[Nic-Ma]

Hi @ericspod ,

Could you please help share some information about this UNet question?

Thanks in advance.

[Minxiangliu]

@Nic-Ma Thank you very much for your quick reply, I want to correct my question a little. How to use the existing model to connect different losses and make different predictions?
EX:

                     |---Dense().....Dense(sigmoid)
monai.networks.nets  |---Dense().....Dense(sigmoid)
                     |---Dense().....Dense(linear)

[ericspod]

I don't think it's a straight forward question of what to do with the existing UNet architecture to do something like this. The network in the blog is composed of two distinct branches and has two outputs, so the equivalent idea here could be to have distinct UNet instances for the tasks you want, and wrap these in a class that broadcasts the inputs to each. This class is what you'd pass to the training engine, and similarly you could have a wrapper around the loss functions you want to use so that they appear as one. Something like this perhaps:

class NetWrapper(nn.module):
    def __init__(self, nets):
        self.nets=nets
    def forward(self, x):
        return tuple(net(x) for net in nets)

class LossWrapper(nn.module):
    def __init__(self,losses):
        self.losses=losses
    def forward(self, inputs, target):
        return sum(loss(input, grnd) for loss, input, grnd in zip(self.losses, inputs, target))

[LaiaHumbertVidan]

Hi,

I was going to ask a similar question. I would like to predict sequential grayscale images (delta HU maps indicating change in image intensity over time with respect to a baseline CT scan) as inputs based on the baseline CT image as input.
Is it possible to train a Unet for regression (to produce grayscale maps) that has multiple outputs, one grayscale map at each time point?

Many thanks!
Laia

[ericspod]

Hi @LaiaHumbertVidan In theory if you have a fixed number of timepoints you could train a UNet or any other autoencoder to produce an output with that many channels. If it's an arbitrary number you may have to train the network in a conditional way with the time point as a condition. I'm not sure UNet would be best suited for that or not.

[SomeUserName1]

Mulitple output channels for one output head is one option:
Here the input is phase-cycle bSSFP MRI data with 12 phase cyles and one phase and one magnitude channel per phase-cycle, resulting in 24 input channels.
The output is 6 channels DWI Tensor with D_xx, D_xy, D_xz, Dyy, D_yz, D_zz as channels.

The follwoing UNet will at least train on the image to image translation task:

unet = monai.networks.nets.UNet(
                spatial_dims=3,
                in_channels=24,
                out_channels=6,
                channels=(24, 48, 96, 196, 384),
                strides=(2, 2, 2, 2),
                dropout=0.1,
                num_res_units=2,
                )

Another option is to add a head per output:
In the following example I add input heads instead of output heads (for pre-training a DWI Tensor auto-encoder, then fine-tuning the image to image translation):

class PreTrainUnet(torch.nn.Module):
    def __init__(self, state):
        super().__init__()
        self.state = state
        self.dwi_tensor_input = monai.blocks.RegistrationResidualConvBlock(
                spatial_dims=3, in_channels=6, out_channels=24, num_layers=10)
        self.bssfp_input = monai.blocks.RegistrationResidualConvBlock(
                spatial_dims=3, in_channels=24, out_channels=24, num_layers=10)

        self.strides = (2, 2, 2, 2)
        self.unet = monai.networks.nets.UNet(
                spatial_dims=3,
                in_channels=24,
                out_channels=6,
                channels=(24, 48, 96, 196, 384),
                strides=self.strides,
                dropout=0.1,
                num_res_units=2,
                )
		# add layers as class vars, s.t. self.parameters() finds everything
        self.all_layers = list(self.unet.children()) + list(self.dwi_tensor_input.children()) + list(self.bssfp_input.children())

        for i, layer in enumerate(self.all_layers):
            setattr(self, f'{layer.__class__.__name__}_{i}', layer)

    def setup(self, checkpoint):
        if self.state == TrainingState.PRETRAIN:
            \# Make model autoencoder with all layers trainable
            for layer in self.all_layers:
                for param in layer.parameters():
                    param.requires_grad = True

        elif self.state == TrainingState.TRANSFER:
            # Freeze all layers but the bssfp input head
            for layer in self.all_layers:
                for param in layer.parameters():
                    param.requires_grad = False
            self.bssfp_input.requires_grad = True

        elif self.state == TrainingState.FINE_TUNE:
            #   Unfreeze all layers and train with a smaller learning rate which is adapted in the L.LightningModule.
            for layer in self.all_layers:
                for param in layer.parameters():
                    param.requires_grad = True
            self.dwi_tensor_input.requires_grad = False

    def forward(self, x):
        if self.state == TrainingState.PRETRAIN:
            x = self.dwi_tensor_input(x)
        else:
            x = self.bssfp_input(x)

        x = self.unet(x)
        return x

To adapt this for multiple output layers first pass the input through the UNet and then to the specific output heads.
Return the outputs of each head.
Regarding the losses: depending on your setup you have to adapt the training loop such that the respective element of the output vector is used to calculate the corresponding loss. For example with Lightning and self.losses a list containing the desired loss for each output head, the training_step could look like this:

def training_step(self, batch, batch_idx):
	x, y_regr, y_seg, y_label = self.unpack_batch(batch)  # Depends on your data loader and targets
	y_hat_regr, y_hat_seg, y_hat_label = self.net(x)  # Calling forward
	
	losses_step = []
	for loss, y, y_hat in zip(self.losses, [y_regr, y_seg, y_label], [y_hat_regr, y_hat_seg, y_hat_label]):
		losses_step.append(
			loss(y_hat, y)
		)
		
	losses_step = torch.tensor(losses_step)
	# Here you have to decide what to do with your losses.  For example one of the following:
	# Use the maximum? 
	loss = torch.max(losses_step)
	# Average ?
	loss = torch.mean(losses_step)
	# weight the losses with self.weights = torch.tensor([0.3, 0.5, 0.2]) for example
	loss = torch.sum(losses_step * self.weights)
	
	return loss	

If you dont want to accumulate your different losses you'd have to write the training loop yourself. Haven't done that but i guess it may be possible to do something like

	for loss, y, y_hat in zip(self.losses, [y_regr, y_seg, y_label], [y_hat_regr, y_hat_seg, y_hat_label]):
		loss_step =  loss(y_hat, y)
		optimizer.zero_grad()
        loss_step.backward()
        optimizer.step()

or similar.
Reading up on the autograd mechanics amd How to use an optimizer/per-parameter options before doing the advanced version is probably a good starting point.

Related reads:
Review on MTL in DNNs by S. Ruder, 2017