I have got a difficult network to construct and I am not sure how to do it right.
I have got an input part and another part consisting of sequentially added blocks of the same type.
I draw a picture of it that you can understand it easily:

diagram.png962×321 7.55 KB


How to I declare the output of one layer as the input of another one in the constructor of my network?
Where should I add my parameters as class members to be updated correctly?

Here is a minimum working example:
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data_utils
import torch.optim as optim
import torch.nn.init as init
import torch.autograd as AG

# Define Custom Layer
class CLayer(AG.Function):
    def __init__(self, tmp_):
        super(CLayer, self).__init__()
        self.tmp = tmp_
    
    def forward(self):
        output = self.tmp + torch.ones(self.tmp.size())
        return output


# Define Block    
class Block(nn.Module):
    def __init__(self, A_, x_, y_):
        super(Block, self).__init__()
        self.A = A_
        self.x = x_
        self.y = y_
        self.customLayer = CLayer(self.x)
        
    def forward(self):
        tmp = self.x - torch.transpose(self.A)*self.y
        output = CLayer(tmp)
        return output
    
    
# Define Network
class Net(nn.Module):
    def __init__(self, A_, x0_):
        super(Net, self).__init__()
        self.x0 = x0_
        self.A = A_
        self.fc1 = nn.Linear(50, 50)
        self.fc1.weight.data = self.A
        self.block0 = Block(self.A, x0, ??)
        self.blockN = nn.ModuleList([Block(self.A, ??, ??) for i in range(2)])
        
    def forward(self):
        y = self.fc1(x)   
        x = self.block0(self.A, self.x0, y)
        for i, l in enumerate(self.ihts):
            x = self.blockN(self.A, x, y)    
        return x
        
        
A_        = torch.randn(m,n)
x0_       = torch.zeros(1,n)
model     = Net(A_, x0_)