Create a dictionnary of layers in pytorch

Hi everyone,

I’m working on an electromagnetism project using neural network and I’m trying to create a nn.Module class for a neural network for which many layers are contained in a dictionnary. Like this:

class ComplexNet(nn.Module):

def __init__(self):
    
    super(ComplexNet, self).__init__()
    
    self.R = dict() 

    for pp in ['00', '01', '10', '11']:
        for ii in ['x','y','z']:
            for jj in ['x','y','z']:  
                for kk in ['r', 'i']:
                    self.R[pp+'e'+jj+'e'+ii+'_'+kk] = HadamardProduct(Nky, Nkx, pp+'e'+jj+'e'+ii+'_'+kk, bias = False)
                    self.R[pp+'h'+jj+'e'+ii+'_'+kk] = HadamardProduct(Nky, Nkx, pp+'h'+jj+'e'+ii+'_'+kk, bias = False)

where HadamardProduct is a layer for dot product operation:

class HadamardProduct(nn.Module):

def __init__(self, Nky, Nkx, string, bias=True):
    super().__init__()
    self.Nky =  Nky
    self.Nkx =  Nkx
    self.weight = nn.Parameter(torch.load(f'{path}/R{string}')).double()
    
def forward(self, input):
    return input *  self.weight

The problem arises when I want to insert model.parameters() in the optimizer. It doesn’t recognize the different layers. Do you have any clue how to fix this issue?
Otherwise, I’ll need to create layers like this:

self.R00exexr = HadamardProduct(Nky, Nkx, ‘00exex_r’, bias = False)
self.R00exeyr = HadamardProduct(Nky, Nkx, ‘00exey_r’, bias = False)
self.R00exezr = HadamardProduct(Nky, Nkx, ‘00exez_r’, bias = False)

But there are 4x3x6 different ones.

Any help would be greatly appreciated!

Victor

Ok, I think I found a trick using nn.ModuleDict

class ComplexNet(nn.Module):

def __init__(self):
    
    super(ComplexNet, self).__init__()
    
    self.R = nn.ModuleDict()

    for pp in ['00', '01', '10', '11']:
        for ii in ['x','y','z']:
            for jj in ['x','y','z']:  
                for kk in ['r', 'i']:
                    self.R[pp+'e'+jj+'e'+ii+'_'+kk] = HadamardProduct(Nky, Nkx, pp+'e'+jj+'e'+ii+'_'+kk, bias = False)
                    self.R[pp+'h'+jj+'e'+ii+'_'+kk] = HadamardProduct(Nky, Nkx, pp+'h'+jj+'e'+ii+'_'+kk, bias = False)

Let me know if you have any other suggestion :slight_smile:

Hi,

ModuleDict is the right thing to use here indeed!