Pytorch 中的 torch.nn.Identity( ) 的作用

最新推荐文章于 2025-04-09 03:08:54 发布

怎样才能回到过去

最新推荐文章于 2025-04-09 03:08:54 发布

阅读量7.1k

点赞数 11

CC 4.0 BY-SA版权

分类专栏： Pytorch 中的各种函数文章标签： Pytorch

本文链接：https://blog.csdn.net/Z2572862506/article/details/129136354

Pytorch 中的各种函数专栏收录该内容

54 篇文章

订阅专栏

torch.nn.Identity()作为恒等函数，用于在深度学习模型中保持输入输出不变。在示例中，它替换预训练模型的最后一层，允许提取特征并自定义新的分类层。这在迁移学习场景中常见，以利用预训练模型的特征表示。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1 torch.nn.Identity( ) 的作用

torch.nn.Identity( ) 相当于一个恒等函数

f(x) = x

这个函数相当于输入什么就输出什么, 可以用在对已经设计好模型结构的修改, 比如模型的最后一层是 1000 分类, 我们可以将最后一层用 nn.Identity( ) 替换掉, 得到它之前学习的特征, 然后再自己设计最后一层的结构

在迁移学习中经常使用

2 示例

import torch
from torch import nn
from torch.nn import NLLLoss
import timm


class MiniModel(nn.Module):
    
    def __init__(self, backbone, num_class, pretrained=False, backbone_ckpt=None):
        
        super().__init__()
        self.backbone = timm.crear_model(backbone, pretrained=pretrained, checkpoint_path=backbone_ckpt)
        self.head = nn.Linear(self.backbone.get_classifier().in_features, num_class)
        # 替代最后一层的全连接网络
        self.backbone.head.fc = nn.Identity()
        self.loss_fn = NLLLoss()
        
    def forward(self, image, label):
        embed = self.backbone(image)
        logit = self.head(embed)
        
        if label is not None:
            logit_logsoftmax = torch.log_softmax(logit, 1)
            loss = self.loss_fn(logit_logsoftmax, label)
            return {"prediction": logit, "loss": loss}
        return {"prediction": logit}