Models with Multiple Outputs

What if your model returns two or more output images? In Nuke, only one output image will be returned from the Inference node. Therefore, only one tensor can be returned from your model forward function. If your model returns more than one output image, you will need to join your output images together to create a single image to return.

Torchscript

Since Inference only returns one output image, the model forward() function will also only return one output tensor. Consider an example model who’s forward function returns two RGB tensors, img1 and img2, as output:

def forward(self, input):
   """
   :param input: A torch.Tensor of size 1 x 3 x H x W
   :return img1: A torch.Tensor of size 1 x 3 x H x W
   :return img2: A torch.Tensor of size 1 x 3 x H x W
   """
   [img1, img2] = self.model.forward(input)
   return img1, img2

This forward() function can be redefined as follows to ensure that it returns one tensor only:

def forward(self, input):
   """
   :param input: A torch.Tensor of size 1 x 3 x H x W
   :return: A torch.Tensor of size 1 x 6 x H x W
   """
   [img1, img2] = self.model.forward(input)
   output = torch.cat((img1, img2), 1)
   return output

where output now has 6 channels.

CatFileCreator

Now that our model is defined so that the forward function returns 6 channels, we use the Channels Out knob in the CatFileCreator node to define where we want these 6 channels to appear in the output image. For example, setting Channels Out as:

will ensure that img1 will appear in the rgba.red, rgba.green and rgba.blue channels while img2 will appear in the forward.u, forward.v and backward.u channels.

Note

Any of the default Nuke channels or any channels that are already defined in your script can be used in the Channels In and Channel Outs knobs of the CatFileCreator.

Inference

Having created the .cat file, we can connect it to the input image and see that the output contains the rgba.red, rgba.green, rgba.blue, forward.u, forward.v and backward.u channels. We can use two shuffle nodes after the Inference node to split this output image back up into two RGB images: