打印网络名称以及相应shape

import torch
from torch import nn

net05=nn.Sequential(
    nn.Conv2d(1,96,kernel_size=11,
              stride=4,padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3,stride=2),
    nn.Conv2d(96,256,kernel_size=5,padding=2),
    nn.MaxPool2d(kernel_size=3,stride=2),
    
    nn.Conv2d(256,384,kernel_size=3,padding=1),
    nn.ReLU(),
    nn.Conv2d(384,384,kernel_size=3,padding=1),
    nn.ReLU(),
    nn.Conv2d(384,256,kernel_size=3,padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=3,stride=2),
    nn.Flatten(),
    nn.Linear(6400,4096),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(4096,4096),
    nn.ReLU(),
    nn.Dropout(p=0.5),
    nn.Linear(4096,10)
)

XX00=torch.randn(1,1,224,224)
for layer in net05:
    XX00=layer(XX00)
    print(layer.__class__.__name__,"output shape:\t",XX00.shape)

你可能感兴趣的:(python,Pytorch,API,python,pytorch,深度学习)