Dynamic decision about execution of a layer

1

Hi.

Let us say I have 10 layers in a neural network. If the input is of type A, I want to execute all 10 layers. If the input if of type B, I want to execute only 8 layers and take the output. Is this easily implementable in pytorch? I tried this in tensorflow and caffe and faced some issues.

Thank you
Athindran

2

Yes, you just use an if statement in your forward pass:

# Create self.layers as a nn.ModuleList
def forward(self, input, input_type):
    out = input
    if input_type == 0:
      for i in range(10):
        out = self.layers[i](out)
    else:
      for i in range(8):
        out = self.layers[i](out)
    return out
3

Follow up,

what if input_type is a control flow. e.g. a tensor: 2 x 5: batch number = 2, each vector contains 5 variables. For example: [1,1,0,0,1]. How to implement this structure?

For example:
[[1,1,0,0,1],[0,1,0,1,0]]

The control flow decide which layer to go, but different sequence has different control flow.