Question

class Model(nn.Module):
  def __init__(self):
    super(Model, self).__init__()
    self.net = nn.Sequential(
      nn.Conv2d(in_channels = 3, out_channels = 16), 
      nn.ReLU(), 
      nn.MaxPool2d(2),
      nn.Conv2d(in_channels = 16, out_channels = 16), 
      nn.ReLU(),
      Flatten(),
      nn.Linear(4096, 64),
      nn.ReLU(),
      nn.Linear(64, 10))

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

我在没有神经网络知识的情况下创建了这个模型，我只是固定参数，直到在训练中起作用为止。我不确定如何获取每一层的输出尺寸（例如，第一层之后的输出尺寸）。

在Pytorch中有一种简单的方法吗？

Answer 1

一种简单的方法是：

将输入传递给模型。
通过每层后打印输出的大小。

class Model(nn.Module):
  def __init__(self):
    super(Model, self).__init__()
    self.net = nn.Sequential(
      nn.Conv2d(in_channels = 3, out_channels = 16), 
      nn.ReLU(), 
      nn.MaxPool2d(2),
      nn.Conv2d(in_channels = 16, out_channels = 16), 
      nn.ReLU(),
      Flatten(),
      nn.Linear(4096, 64),
      nn.ReLU(),
      nn.Linear(64, 10))

  def forward(self, x):
    for layer in self.net:
        x = layer(x)
        print(x.size())
    return x

model = Model()
x = torch.randn(1, 3, 224, 224)

# Let's print it
model(x)

但是请注意输入大小，因为您在网络中使用nn.Linear。如果您的输入大小不是4096，则会导致nn.Linear的输入大小不兼容。

Answer 2

就像吴伟达的答案，但要短一些：

def get_output_shape(model, image_dim):
    return model(torch.rand(*(image_dim))).data.shape

在此示例中，我需要找出最后一个线性层的输入：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.expected_input_shape = (1, 1, 192, 168)
        self.conv1 = nn.Conv2d(1, 32, 3, 1) 
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.maxpool1 = nn.MaxPool2d(2)
        self.maxpool2 = nn.MaxPool2d(3)

        # Calculate the input of the Linear layer
        conv1_out = get_output_shape(self.maxpool1, get_output_shape(conv1, self.expected_input_shape))
        conv2_out = get_output_shape(self.maxpool2, get_output_shape(conv2, conv1_out)) 
        fc1_in = np.prod(list(conv2_out)) # Flatten

        self.fc1 = nn.Linear(fc1_in, 38)

    def forward(self, x):
        x = self.conv1(x) 
        x = F.relu(x)
        x = self.maxpool1(x) 
        x = self.conv2(x)
        x = F.relu(x)
        x = self.maxpool2(x) 
        x = self.dropout1(x) 
        x = torch.flatten(x, 1) # flatten to a single dimension
        x = self.fc1(x) 
        output = F.log_softmax(x, dim=1) 
        return output

这样，如果我对以前的图层进行了更改，则无需重新计算！

我的答案基于this answer

Answer 3

这是一个辅助函数形式的解决方案：

def get_tensor_dimensions_impl(model, layer, image_size, for_input=False):
    t_dims = None
    def _local_hook(_, _input, _output):
        nonlocal t_dims
        t_dims = _input[0].size() if for_input else _output.size()
        return _output    
    layer.register_forward_hook(_local_hook)
    dummy_var = torch.zeros(1, 3, image_size, image_size)
    model(dummy_var)
    return t_dims

示例：

from torchvision import models, transforms

a_model = models.squeezenet1_0(pretrained=True) 
get_tensor_dimensions_impl(a_model, a_model._modules['classifier'], 224)

输出为：

<块引用>

torch.Size([1, 1000, 1, 1])

Answer 4

例如，您可以将Torchsummary用于ImageNet尺寸（3x224x224）：

import time

from concurrent.futures import ThreadPoolExecutor

def f(t):
    time.sleep(t)
    return t

pool = ThreadPoolExecutor(max_workers=300)
res = list(pool.map(f, [2, 10], timeout=1))

来源：model-summary-in-pytorch

Answer 5

获取nn.Sequential容器中特定层之后的大小的另一种方法是添加一个自定义Module，它仅打印出输入的大小。

class PrintSize(nn.Module):
  def __init__(self):
    super(PrintSize, self).__init__()
    
  def forward(self, x):
    print(x.shape)
    return x

现在您可以这样做：

model = nn.Sequential(
    nn.Conv2d(3, 10, 5, 1),
    // lots of convolutions, pooling, etc.
    nn.Flatten(),
    PrintSize(),
    nn.Linear(1, 12), // the input dim of 1 is just a placeholder
)

现在，您可以执行model(x)，并且在运行Conv2d层后，它将打印输出的形状。如果您有很多卷积并且想弄清楚第一个完全连接的层的最终尺寸是什么，这将很有用。您无需将nn.Sequential重新格式化为模块，只需一行就可以进入此帮助程序类。

Answer 6

也许你可以试试print(model.state_dict()['next_layer.weight'].shape)。这为您提供了最后一层的输出形状的提示。

如何在Pytorch中获得神经网络每一层的输出尺寸？

6 个答案: