In-Place操作及onnx导出

本文介绍了PyTorch中的In-Place操作,这种操作能节省GPU显存但可能导致梯度计算问题。文章提醒在使用In-Place操作时要谨慎,特别是在导出ONNX模型时应避免,因为这可能引入scatterND等不被广泛支持的算子。建议使用非In-Place操作以保证模型移植性和兼容性。

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

本文只做本人笔记记录所用!

一、概念介绍

In-place 操作:被称为就地操作和原地操作

pytorch中原地操作的后缀为_,如.add_()或.scatter_(),

就地操作:是直接更改给定Tensor的内容而不进行复制的操作,即不会为变量分配新的内存(不分配新的对象)。Python操作类似+=或*=也是就地操作。

当然有In-Place就有out-of-place与之对应。

  • 每个 in-place 操作实际上都需要重写计算图的实现。out-of-place只是分配新对象并保留对旧计算图的引用,而 in-place 操作则需要将所有输入的创建更改为代表此操作的函数。

二、介绍

深度学习降低显存使用量,有几种简单的方法可以减少模型占用的GPU显存

(1)更改模型的结构,比如做轻量化模型的设计,结构重参等方法,可以有效降低模型参数量

(2)模型蒸馏,从教师模型(参数量大)-》蒸馏出学生模型(参数量小)

(3)每次训练batchsize设置小一些,也可以减少显存占用量(缺点是模型需要更长的时间训练,且模型性能可能会受损)

在我们pytorch训练

### DenseNet-12 模型的相关信息 DenseNet 是一种经典的卷积神经网络架构,其核心思想在于通过特征重用提高参数效率并缓解梯度消失问题。然而,在常见的公开资源中,通常提到的是 **DenseNet-121**, **DenseNet-169**, 和 **DenseNet-201** 等变体[^5]。关于 **DenseNet-12** 的具体实现或预训练模型文件较少见于主流开源平台。 如果需要获取 DenseNet-12 的 ONNX 文件,可以通过以下方法: #### 方法一:基于 PyTorch 或 TensorFlow 实现自定义导出 可以利用 PyTorch 提供的 `torch.onnx.export` 函数来生成 ONNX 文件。以下是具体的代码示例: ```python import torch from torchvision import models # 定义 DenseNet-12 结构 (假设为简化版) class DenseNet12(torch.nn.Module): def __init__(self, num_classes=1000): super(DenseNet12, self).__init__() # 自定义层结构 from torchvision.models.densenet import _Transition, _DenseBlock, DenseNet self.features = torch.nn.Sequential( torch.nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3), torch.nn.ReLU(inplace=True), torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1), _DenseBlock(num_layers=2, num_input_features=64, bn_size=4, growth_rate=32), _Transition(num_input_features=128, num_output_features=64), torch.nn.AdaptiveAvgPool2d((1, 1)) ) self.classifier = torch.nn.Linear(64, num_classes) def forward(self, x): features = self.features(x) out = torch.flatten(features, 1) out = self.classifier(out) return out model = DenseNet12() dummy_input = torch.randn(1, 3, 224, 224) # 导出ONNX 格式 torch.onnx.export(model, dummy_input, "densenet12.onnx", opset_version=11, input_names=["input"], output_names=["output"]) ``` 上述代码实现了一个简化的 DenseNet 架构,并将其保存为名为 `densenet12.onnx` 的文件。 #### 方法二:从现有模型裁剪得到 DenseNet-12 另一种方式是从已有的 DenseNet-121 预训练模型中提取部分层数构建新的模型。例如,可以从官方提供的 PyTorch 预训练模型加载权重,并仅保留前几层作为基础框架: ```python pretrained_model = models.densenet121(pretrained=True) custom_model = torch.nn.Sequential(*list(pretrained_model.children())[:4]) # 截取特定层 ``` 完成截断后,再按照需求调整剩余部分直至满足目标配置。 --- ### 使用 TensorRT 加速推理过程 一旦获得了 `.onnx` 文件,则可通过 NVIDIA 提供的工具链进一步优化性能表现。例如借助命令行接口执行如下操作即可获得对应的 TensorRT engine 文件: ```bash trtexec --onnx=densenet12.onnx --saveEngine=densenet12.engine --buildOnly --fp16 ``` 这一步骤能够显著提升部署环境中的计算速度与能效比[^4]。 ---
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

曙光_deeplove

你的鼓励是我努力的最大源泉

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值