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