1
$\begingroup$

Does there exist a fast and convenient way for handling such a problem:

class MyModule(nn.Module):
def __init__(self):
    super(MyModule, self).__init__()
    self.linears = nn.ModuleList([nn.Linear(10, 20) for _ in range(10)])

def forward(self, x, indices):
    x = self.linears[indices](x) 
    return x

You see i want to access different layers in the network conditioned on an additional input, which is also a list. Further i want to process the whole batch at once and the output.shape != input.shape.

$\endgroup$
0
$\begingroup$

Here is my understanding of your problem:

# Import
from torch import nn

# Define custom class
class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 20) for _ in range(10)])

    def forward(self, x, indices):
        x = self.linears[indices](x) 
        return x

# Intialize custom class    
net = MyModule()

# Access networks layers based on additional input
additional_input = 1
if additional_input == 1:
    idx = 0
    print(net.linears[idx].in_features)
$\endgroup$
  • $\begingroup$ thank you brian. i think i didnt specify my problem good enough. by accessing the layers i actually meant to make a forward pass for a whole batch at once, while i need different layers for each entry in the batch conditioned on the extra input $\endgroup$ – Andreas Look Mar 1 at 18:57

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.