Defining model parameters

Hi all. I am very new to Pytorch and trying to build a simple spectral GNN by myself.
Following is the code I wrote, but somehow it seems like the model parameters are not properly defined, when I create a network and do the model.parameters(). I don’t know how to fix this…
I want to set W_0 and W_1 as the model parameters…

class Spectral_GNN(nn.Module):
def init(self, input_units, hidden_units, output_units):
super(Spectral_GNN, self).init()

    def normal(shape):
        return torch.randn(size=shape)*0.01
    
    def two():
        return (normal((input_units, hidden_units)),
                normal((hidden_units, output_units)))
    
    self.W0, self.W1 = two()
    
    # Attach gradients
    params = [self.W0, self.W1]
    for param in params:
        param.requires_grad_(True)        
    
def forward(self, X, A):
    D = torch.sum(A, 1)
    D_inv = D**-0.5
    D_inv = torch.diag(D_inv)
    A_hat = D_inv * A * D_inv
    
    aggregate = torch.matmul(A_hat, X)
    propagate = torch.relu(torch.matmul(aggregate, self.W0))
    aggregate = torch.matmul(A_hat, propagate)
    propagate = torch.matmul(aggregate, self.W1)
    
    return propagate

Greetings, you can follow this prescription.

For example simply wrap this in a nn.Parameter:

def normal(shape):
    return nn.Parameter(torch.randn(size=shape)*0.01)
for foo in gnn.parameters():
    print(foo)

Output:
Parameter containing:
tensor([[-0.0022]], requires_grad=True)
Parameter containing:
tensor([[0.0037]], requires_grad=True)
1 Like