在PyTorch中关于形状的推理

Created On: Mar 27, 2023 | Last Updated: Mar 27, 2023 | Last Verified: Not Verified

使用 PyTorch 编写模型时，通常某一层的参数依赖于上一层输出的形状。例如，nn.Linear 层的 in_features 必须与输入数据的 size(-1) 匹配。对于某些层，形状计算涉及复杂公式，例如卷积操作。

一种解决方法是使用随机输入执行前向传播，但这在内存和计算方面是低效的。

相反，我们可以利用 meta 设备来确定层的输出形状，而无需生成任何数据。

import torch
import timeit

t = torch.rand(2, 3, 10, 10, device="meta")
conv = torch.nn.Conv2d(3, 5, 2, device="meta")
start = timeit.default_timer()
out = conv(t)
end = timeit.default_timer()

print(out)
print(f"Time taken: {end-start}")

由于数据未生成，传递任意大的输入数据不会显著影响形状计算所需的时间。

t_large = torch.rand(2**10, 3, 2**16, 2**16, device="meta")
start = timeit.default_timer()
out = conv(t_large)
end = timeit.default_timer()

print(out)
print(f"Time taken: {end-start}")

考虑以下任意定义的网络：

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

通过为每一层注册一个前向钩子，我们可以查看整个网络中每层的中间形状，该钩子会打印输出形状。

def fw_hook(module, input, output):
    print(f"Shape of output to {module} is {output.shape}.")


# Any tensor created within this torch.device context manager will be
# on the meta device.
with torch.device("meta"):
    net = Net()
    inp = torch.randn((1024, 3, 32, 32))

for name, layer in net.named_modules():
    layer.register_forward_hook(fw_hook)

out = net(inp)