Models with Multiple Inputs

What if your model expects two or more input images? In Nuke, only one input image can be passed to the Inference node. Therefore, only one tensor will be passed to your model forward function. If your model requires more than one input image, you will need to join your input images together to create a single image for processing.

Torchscript

Since Inference only accepts one input image, the model forward function will also only have one input tensor. Consider an example model who’s forward() function initially required two tensors, img1 and img2, as input, with each tensor representing an RGB image:

def forward(self, img1, img2):
   """
   :param img1: A torch.Tensor of size 1 x 3 x H x W representing image 1
   :param img2: A torch.Tensor of size 1 x 3 x H x W representing image 2
   :return: A torch.Tensor of size 1 x 1 x H x W of zeros or ones
   """
   output = self.model.forward(img1, img2)
   return output

In order to reformat this so that the forward() function accepts one input tensor only, we can redefine our forward function as follows:

def forward(self, input):
   """
   :param input: A torch.Tensor of size 1 x 6 x H x W representing two RGB images
   :return: A torch.Tensor of size 1 x 1 x H x W of zeros or ones
   """

   # Split the input tensor into two tensors,
   # one representing each image
   input_split = torch.split(input, 3, 1)
   img1 = input_split[0]
   img2 = input_split[1]

   output = self.model.forward(img1, img2)
   return output

where input now has 6 channels and can be split into two tensors, img1 and img2, before passing them to model.forward().

Channel Shuffling

Since the Inference node only accepts one input image, channels from images that you would like processed by the model will have to be shuffled into a single image in Nuke. Since our model forward function now accepts a 6 channel tensor, we can use a shuffle node in Nuke to shuffle two RGB images into a single image with 6 channels - rgba.red, rgba.green, rgba.blue, forward.u, forward.v, and backward.u, for example - and pass this image as input to the Inference node.

CatFileCreator

When converting this model from TorchScript to a .cat file, the Channels In knob should define the 6 channels from the input image to Inference that you would like to process by the model. In our case, the 6 channels that we would like to process are the rgba.red, rgba.green, rgba.blue, forward.u, forward.v, and backward.u channels of the input image, so we define the CatFileCreator knobs as follows:

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. See the appendix for the full list of Nuke channel names that can be used with these knobs.

Inference

When we have created the Inference node, we can see that the 6 channels from the input image that will be processed by the model are listed in Channels In:

This indicates that we need to pass an image to this Inference node that has all 6 of these channels. We can connect this Inference node to the 6 channel image we created earlier to get the final output.