I have a model with a forward pass which shares the same computations in training and inference mode. I don’t want to use a boolean training or something like that to declare if I am currently training. I want a forward pass with two outputs. The first one being the training output and the second one the inference output. Depending on how I call the forward pass / which of the two outputs I pick, I want to compute only the part of the graph actually needed for the output. Sometimes I want to calculate both, training and inference output with one forward pass. I want to have a clean model being able to cope with all three situations efficiently. Below there is a dummy model.
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv2d_1 = torch.nn.Conv2d(3, 3, 3)
self.conv2d_2 = torch.nn.Conv2d(3, 3, 3)
self.conv2d_3 = torch.nn.Conv2d(3, 10, 3)
self.conv2d_4 = torch.nn.Conv2d(10, 10, 3)
self.conv2d_5 = torch.nn.Conv2d(10, 10, 3)
def forward(self, input):
train_out = self.conv2d_1(input)
inference_out = self.conv2d_2(input)
train_out = self.conv2d_3(train_out)
inference_out = self.conv2d_3(inference_out)
train_out = self.conv2d_4(train_out)
inference_out = self.conv2d_5(inference_out)
return train_out, inference_out
if __name__ == '__main__':
model = Model()
training_input = torch.rand(1, 3, 100, 100)
inference_input = torch.rand(1, 3, 200, 200)
training_out, _ = model(training_input)
_ , inference_out = model(inference_input)
training_inference_input = torch.rand(1, 3, 100, 100)
training_out, inference_out = model(training_inference_input)
In this example they share the convolution self.conv2d_3, in reality the model is much more complex. If I call the model like this, still the whole graph gets computed even if I only take one of the outputs. What is the best way to achieve only the part of graph being computed which is needed in the output?