nn.Identity 单位矩阵，同一矩阵

取个名字真难呐2025-03-06 9:30

文章目录

[1. 说明](#1. 说明)
[2. pytorch 代码](#2. pytorch 代码)

1. 说明

在搭建网络结构中，为了保证搭建的网络具有高度扩展性和后续调试模型框架，在保证整体结构完整情况下，用nn.Identity 进行占位符处理。

2. pytorch 代码

pytorch代码

python 复制代码

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, use_dropout=True):
        super(MyModel, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            # 根据配置决定使用 Dropout 还是 Identity
            nn.Identity() if not use_dropout else nn.Dropout(p=0.5),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
        )
        self.classifier = nn.Linear(32 * 32 * 32, 10)  # 假设输入图像尺寸为 32x32

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# 实例化模型时选择是否使用 Dropout
model_with_dropout = MyModel(use_dropout=True)
model_without_dropout = MyModel(use_dropout=False)

# 测试模型输出形状
x = torch.randn(4, 3, 32, 32)
print("Output with dropout:", model_with_dropout(x).shape)
print("Output without dropout:", model_without_dropout(x).shape)

result

python 复制代码

Output with dropout: torch.Size([4, 10])
Output without dropout: torch.Size([4, 10])