Question

我正在尝试在pytorch中创建一个多层神经网络类。我想知道以下两段代码是否创建相同的网络。

具有nn.Linear的模型1

class TestModel(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
    super(TestModel, self).__init__()
    self.fc1 = nn.Linear(input_dim,hidden_dim)
    self.fc2 = nn.Linear(hidden_dim,output_dim)

def forward(self, x):
    x = nn.functional.relu(self.fc1(x))
    x = nn.functional.softmax(self.fc2(x))
    return x

具有nn.Sequential的模型2

class TestModel2(nn.Module):
def __init__(self, input, hidden, output):
    super(TestModel2, self).__init__()
    self.seq = nn.Sequential(
               nn.Linear(input_dim,hidden_dim),
               nn.ReLU(),
               nn.Linear(hidden_dim,output_dim),
               nn.Softmax()
               )

def forward(self, x):
    return self.seq(x)

Answer 1

是的，这两段代码创建了相同的网络。一种使自己确信这是事实的方法是将两个模型都保存到ONNX。

import torch.nn as nn

class TestModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(TestModel, self).__init__()
        self.fc1 = nn.Linear(input_dim,hidden_dim)
        self.fc2 = nn.Linear(hidden_dim,output_dim)

    def forward(self, x):
        x = nn.functional.relu(self.fc1(x))
        x = nn.functional.softmax(self.fc2(x))
        return x   


class TestModel2(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(TestModel2, self).__init__()
        self.seq = nn.Sequential(
                nn.Linear(input_dim, hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, output_dim),
                nn.Softmax()
                )

    def forward(self, x):
        return self.seq(x)

m = TestModel(1, 2, 3)
m2 = TestModel2(1, 2, 3)

torch.onnx.export(m, torch.Tensor([0]), "test.onnx", verbose=True)
/opt/anaconda3/envs/py36/bin/ipython:9: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
graph(%0 : Float(1)
      %1 : Float(2, 1)
      %2 : Float(2)
      %3 : Float(3, 2)
      %4 : Float(3)) {
  %5 : Float(1!, 2) = onnx::Transpose[perm=[1, 0]](%1), scope: TestModel/Linear[fc1]
  %6 : Float(2) = onnx::MatMul(%0, %5), scope: TestModel/Linear[fc1]
  %7 : Float(2) = onnx::Add(%6, %2), scope: TestModel/Linear[fc1]
  %8 : Float(2) = onnx::Relu(%7), scope: TestModel
  %9 : Float(2!, 3!) = onnx::Transpose[perm=[1, 0]](%3), scope: TestModel/Linear[fc2]
  %10 : Float(3) = onnx::MatMul(%8, %9), scope: TestModel/Linear[fc2]
  %11 : Float(3) = onnx::Add(%10, %4), scope: TestModel/Linear[fc2]
  %12 : Float(3) = onnx::Softmax[axis=0](%11), scope: TestModel
  return (%12);
}

torch.onnx.export(m2, torch.Tensor([0]), "test.onnx", verbose=True)
/opt/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py:475: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
  result = self._slow_forward(*input, **kwargs)
graph(%0 : Float(1)
      %1 : Float(2, 1)
      %2 : Float(2)
      %3 : Float(3, 2)
      %4 : Float(3)) {
  %5 : Float(1!, 2) = onnx::Transpose[perm=[1, 0]](%1), scope: TestModel2/Sequential[seq]/Linear[0]
  %6 : Float(2) = onnx::MatMul(%0, %5), scope: TestModel2/Sequential[seq]/Linear[0]
  %7 : Float(2) = onnx::Add(%6, %2), scope: TestModel2/Sequential[seq]/Linear[0]
  %8 : Float(2) = onnx::Relu(%7), scope: TestModel2/Sequential[seq]/ReLU[1]
  %9 : Float(2!, 3!) = onnx::Transpose[perm=[1, 0]](%3), scope: TestModel2/Sequential[seq]/Linear[2]
  %10 : Float(3) = onnx::MatMul(%8, %9), scope: TestModel2/Sequential[seq]/Linear[2]
  %11 : Float(3) = onnx::Add(%10, %4), scope: TestModel2/Sequential[seq]/Linear[2]
  %12 : Float(3) = onnx::Softmax[axis=0](%11), scope: TestModel2/Sequential[seq]/Softmax[3]
  return (%12);
}

因此，两个模型都将以相同的操作生成相同的ONNX图。

torch.nn.sequential与多个torch.nn.linear的组合

1 个答案: