Как определить размеры входных данных PyTorch для визуализации модели?

Question

Я переезжаю из Keras в PyTorch и с трудом создаю графическую визуализацию моделей. В Keras вы можете просто вызвать plot_model(), и эта функция будет использовать graphviz для создания хорошего графика.

Однако в PyTorch это не так просто. Я видел, что есть много доступных инструментов, таких как TorchViz, но все они требуют, чтобы вы передали некоторый вклад в модель . Например:

1. Этот вопрос StackOverflow ( Как визуализировать net в Pytorch? ) предлагает использовать пакет torchviz следующим образом:

x = torch.zeros(1, 3, 224, 224, dtype=torch.float, requires_grad=False)
out = resnet(x)
make_dot(out) 

Пакет HiddenLayer ( пример ) также требует фиктивного ввода:

model = torchvision.models.vgg16()
g = hl.build_graph(model, torch.zeros([1, 3, 224, 224]))

Поэтому мой вопрос: Предположим, я скачал предварительно обученная модель, которая имеет мало документации. Как определить, какой размер фиктивного ввода необходим для запуска этих функций визуализации?

Например, когда я print() модель Ale xNet, какая часть приведенного ниже результата говорит мне, что ввод должен быть [1, 3, 224, 224]?

model = torchvision.models.alexnet(pretrained=True)
print(model)

# AlexNet(
#   (features): Sequential(
#     (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))
#     (1): ReLU(inplace=True)
#     (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
#     (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
#     (4): ReLU(inplace=True)
#     (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
#     (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#     (7): ReLU(inplace=True)
#     (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#     (9): ReLU(inplace=True)
#     (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#     (11): ReLU(inplace=True)
#     (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
#   )
#   (avgpool): AdaptiveAvgPool2d(output_size=(6, 6))
#   (classifier): Sequential(
#     (0): Dropout(p=0.5, inplace=False)
#     (1): Linear(in_features=9216, out_features=4096, bias=True)
#     (2): ReLU(inplace=True)
#     (3): Dropout(p=0.5, inplace=False)
#     (4): Linear(in_features=4096, out_features=4096, bias=True)
#     (5): ReLU(inplace=True)
#     (6): Linear(in_features=4096, out_features=1000, bias=True)
#   )
# )

Первый слой говорит: Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2)). Как я узнаю из этого факта, что мне нужно создать [1, 3, 224, 224] ввод?