如何查看PyTorch预训练模型及model_ir_se50.pth的网络架构？

Hey there! Let's tackle your two PyTorch model architecture questions—they're super common when working with pretrained models, so I've got you covered.

There are two straightforward methods to inspect a pretrained model's structure, depending on how much detail you need:

• Method 1: Print the Model Directly
  This is the simplest way to get a quick overview of all layers and their configurations. Here's an example using a ResNet50 model from torchvision:

  import torch
  import torchvision.models as models
  
  # Load the pretrained model
  model = models.resnet50(pretrained=True)
  # Print the full architecture
  print(model)
  

  The output will list every layer (convolutional, pooling, fully connected, etc.) along with their parameters like kernel size, stride, and number of filters—perfect for a high-level look at the model's layout.

• Method 2: Use torchsummary for Detailed Layer Information
  If you want more structured data (like input/output dimensions, parameter counts, and total trainable parameters), torchsummary is your go-to tool. First install it via pip:

  pip install torchsummary
  

  Then use it with your model (note: you need to specify the input size expected by the model):

  from torchsummary import summary
  
  # For a standard image model, input size is (channels, height, width)
  summary(model, input_size=(3, 224, 224))
  

  This will output a table with each layer's type, output shape, parameter count, and a summary of total parameters (trainable vs. non-trainable).

PyTorch .pth files only store model weights, not the actual network structure. So first, you need the model class definition that matches the weights. Here's how to do this:

• Step 1: Get the IR-SE50 Model Class Definition
  The IR-SE50 model is part of the InsightFace family of face recognition models. You'll need the corresponding model class (you can find this in the original open-source repo, or use this standard implementation):

  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  class Bottleneck(nn.Module):
      def __init__(self, in_channels, out_channels, stride, use_se):
          super().__init__()
          self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
          self.bn1 = nn.BatchNorm2d(out_channels)
          self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
          self.bn2 = nn.BatchNorm2d(out_channels)
          self.conv3 = nn.Conv2d(out_channels, out_channels*4, kernel_size=1, bias=False)
          self.bn3 = nn.BatchNorm2d(out_channels*4)
          self.se = SELayer(out_channels*4) if use_se else nn.Identity()
          self.shortcut = nn.Sequential()
          if stride != 1 or in_channels != out_channels*4:
              self.shortcut = nn.Sequential(
                  nn.Conv2d(in_channels, out_channels*4, kernel_size=1, stride=stride, bias=False),
                  nn.BatchNorm2d(out_channels*4)
              )
  
      def forward(self, x):
          out = F.relu(self.bn1(self.conv1(x)))
          out = F.relu(self.bn2(self.conv2(out)))
          out = self.bn3(self.conv3(out))
          out = self.se(out)
          out += self.shortcut(x)
          return F.relu(out)
  
  class SELayer(nn.Module):
      def __init__(self, channel, reduction=16):
          super().__init__()
          self.avg_pool = nn.AdaptiveAvgPool2d(1)
          self.fc = nn.Sequential(
              nn.Linear(channel, channel//reduction, bias=False),
              nn.ReLU(inplace=True),
              nn.Linear(channel//reduction, channel, bias=False),
              nn.Sigmoid()
          )
  
      def forward(self, x):
          b, c, _, _ = x.size()
          y = self.avg_pool(x).view(b, c)
          y = self.fc(y).view(b, c, 1, 1)
          return x * y.expand_as(x)
  
  class IR_SE50(nn.Module):
      def __init__(self, num_classes=512):
          super().__init__()
          self.input_layer = nn.Sequential(
              nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False),
              nn.BatchNorm2d(64),
              nn.PReLU(64)
          )
          self.layer1 = self._make_layer(64, 64, 3, stride=1, use_se=True)
          self.layer2 = self._make_layer(256, 128, 4, stride=2, use_se=True)
          self.layer3 = self._make_layer(512, 256, 14, stride=2, use_se=True)
          self.layer4 = self._make_layer(1024, 512, 3, stride=2, use_se=True)
          self.output_layer = nn.Sequential(
              nn.BatchNorm2d(2048),
              nn.Dropout(0.4),
              nn.Flatten(),
              nn.Linear(2048, num_classes),
              nn.BatchNorm1d(num_classes)
          )
          self._initialize_weights()
  
      def _make_layer(self, in_channels, out_channels, blocks, stride, use_se):
          layers = [Bottleneck(in_channels, out_channels, stride, use_se)]
          for _ in range(1, blocks):
              layers.append(Bottleneck(out_channels*4, out_channels, 1, use_se))
          return nn.Sequential(*layers)
  
      def _initialize_weights(self):
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
              elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
              elif isinstance(m, nn.Linear):
                  nn.init.normal_(m.weight, 0, 0.01)
                  nn.init.constant_(m.bias, 0)
  
      def forward(self, x):
          x = self.input_layer(x)
          x = self.layer1(x)
          x = self.layer2(x)
          x = self.layer3(x)
          x = self.layer4(x)
          return self.output_layer(x)
  

• Step 2: Load the Weights and Inspect the Architecture
  Now instantiate the model, load the .pth weights, and print the structure (or use torchsummary for more details):

  # Create an instance of the IR_SE50 model
  model = IR_SE50()
  # Load the pretrained weights (make sure the path to model_ir_se50.pth is correct)
  model.load_state_dict(torch.load('model_ir_se50.pth'))
  # Print the full architecture
  print(model)
  
  # Optional: Use torchsummary to see input/output dimensions and parameter counts
  from torchsummary import summary
  # IR-SE50 typically expects 3-channel 112x112 face images
  summary(model, input_size=(3, 112, 112))
  

  If you get a "key mismatch" error when loading weights, double-check that your model class matches the original implementation (the parameter names need to align exactly with those in the .pth file).

内容的提问来源于stack exchange，提问作者kritika Garg