libtorch释放显存:torch::NoGradGuard no_grad;

一般使用.reset()就够了

但是神经网络的output是释放不掉的,因为系统要计算梯度。

怎么办?

torch::NoGradGuard no_grad;

torch::NoGradGuard 的主要作用就是在其作用域内禁用自动梯度计算,从而节省显存和计算资源。变量 no_grad 本身不需要被直接使用,它的作用是通过作用域来实现的。

示例代码如下

#include <torch/torch.h>
#include <iostream>
#include <vector>

// 假设你已经定义了一个模型 m_module
struct MyModule : torch::nn::Module {
    torch::nn::Linear fc{nullptr};

    MyModule() {
        fc = register_module("fc", torch::nn::Linear(10, 2));
    }

    torch::Tensor forward(torch::Tensor x) {
        return fc->forward(x);
    }
};

int main() {
    // 创建一个模型实例并将其移动到GPU(如果可用)
    MyModule m_module;
    m_module.to(torch::kCUDA);

    // 创建一个输入张量并将其移动到GPU(如果可用)
    torch::Tensor input = torch::randn({1, 10}).to(torch::kCUDA);

    // 创建一个向量来存储输出
    std::vector<torch::Tensor> outputs;

    // 使用 torch::NoGradGuard 确保在以下范围内不计算梯度
    {
        torch::NoGradGuard no_grad;

        // 前向传播
        torch::Tensor output = m_module.forward(input).toTensor();
        
        // 将输出转移到 CPU 上
        auto output1 = output.to(torch::kCPU);
        
        // 清空原始的输出 tensor
        output.reset();
        
        // 将结果存储在输出向量中
        outputs.push_back(output1);
    } // 这里 no_grad 作用域结束,梯度计算自动恢复

    // 输出结果
    for (const auto& out : outputs) {
        std::cout << out << std::endl;
    }

    return 0;
}

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值