一般使用.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;
}

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐