训练在另一个模型（Pytorch）中实例化的模型

Question

我有两个 class 神经元网络，一个是 GNN 类型，另一个是简单的线性类型，后者在第一个中实例化！我怎样才能同时训练两者？这是一个例子：

class linear_NN(nn.Module):
  
  def __init__(self, input_dim, out_dim...):
    super().__init__()

  def forward(self, x, dim = 0):
    '''Forward pass'''
    return x

主要class或大class

class GNN(nn.Module):
  
  def __init__(self, input_dim, n-hidden, out_dim...):
    super().__init__()

  def forward(self, h, dim = 0):
    '''Forward pass'''
    model=linear_NN(input, out..)
    model(h, dim)
    return h

Answer 1

您必须在 __init__(...):

中声明它

class GNN(nn.Module):
  def __init__(self, input_dim, n-hidden, out_dim, ...):
    super().__init__()
    self.linear = linear_NN(input, out..)

  def forward(self, h, dim = 0):
    '''Forward pass'''
    self.linear(h, dim)
    return h

然后，self.linear 模型将注册到您的 GNN 主模型，如果您得到 GNN(...).parameters()，您将在那里看到线性参数。

训练在另一个模型（Pytorch）中实例化的模型

train a model which is instantiated in another model ( Pytorch)

python

neural-network

pytorch