必须继承 torch::nn::Module,用 register_parameter() 注册参数,禁用原地操作,确保设备一致且梯度清零。

如何在 LibTorch 中继承 torch::nn::Module 定义自定义层
必须继承 torch::nn::Module(而非裸写函数),否则无法被 torch::nn::Sequential 管理、不能自动注册参数、也无法参与反向传播。构造函数里要用 register_module() 或 register_parameter() 显式注册所有可学习张量。
常见错误是直接在类成员里声明 torch::Tensor,比如:torch::Tensor weight_; —— 这样不会被 parameters() 返回,梯度根本不会更新。
正确做法示例:
struct MyLinearImpl : torch::nn::Module {
MyLinearImpl(int64_t in, int64_t out) {
weight_ = register_parameter("weight", torch::randn({out, in}));
bias_ = register_parameter("bias", torch::zeros({out}));
}
torch::Tensor forward(torch::Tensor x) {
return torch::matmul(x, weight_.t()) + bias_;
}
torch::Tensor weight_, bias_;
};
TORCH_MODULE(MyLinear);
注意两点:register_parameter() 返回引用并赋值给成员;TORCH_MODULE() 宏必须调用,否则无法使用默认构造和拷贝语义。
为什么 forward 里不能用原地操作(in-place)
LibTorch 的自动微分引擎(autograd)依赖计算图中每个节点的输入输出关系。一旦用了 .add_()、.mul_()、.relu_() 等带下划线的原地方法,会破坏中间变量的生命周期或覆盖原始 tensor 的 grad_fn,导致反向传播时抛出 RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation。
实操建议:
- 一律用非原地版本:用
torch::relu(x),不用x.relu_() - 若必须原地(如内存受限),需在操作前调用
x.retain_grad()并确保不破坏依赖链,但一般不推荐 - 调试时加
torch::set_grad_enabled(true)并检查x.requires_grad()是否为 true
如何手动实现 backward(绕过 autograd)
95% 的场景不需要手写 backward —— 只要 forward 全部由可导的 LibTorch 操作构成,autograd 就能自动推导。只有两种情况才需重载 backward():
- forward 包含不可导操作(如 argmax、硬阈值、自定义 CUDA kernel)
- 需要定制梯度行为(如梯度裁剪、伪量化、直通估计 STE)
此时应继承 torch::autograd::Function,定义 forward() 和 backward() 静态方法,并在 forward 中调用 apply()。例如 STE 层:
struct STEFunction : torch::autograd::Function<stefunction> {
static torch::Tensor forward(torch::autograd::AutogradContext* ctx,
torch::Tensor x) {
auto out = (x > 0).to(torch::kFloat);
ctx->save_for_backward({x});
return out;
}
static std::vector<:tensor> backward(
torch::autograd::AutogradContext* ctx,
std::vector<:tensor> grad_outputs) {
auto saved = ctx->get_saved_variables();
return {grad_outputs[0].clamp(-1, 1)}; // 梯度截断
}
};
// 在 Module 的 forward 中调用:
// return STEFunction::apply(x);
</:tensor></:tensor></stefunction>
注意:手写 backward() 时,返回梯度数量、形状、设备必须与 forward 输入严格一致,否则训练会静默失败或崩溃。
训练时 loss.backward() 后参数没更新?检查这三处
最常被忽略的是优化器未绑定参数、梯度未清零、或参数不在同一设备上。LibTorch 不像 PyTorch Python 版有隐式保护,出错往往无提示。
关键检查点:
-
optimizer.step()前必须调用optimizer.zero_grad()—— C++ 版不会自动清零 - 确认
model->parameters()返回的参数确实包含你注册的 tensor(打印 size 或 names 验证) - 所有 tensor(input、label、model)必须在同一 device 上:
model->to(device)后,数据也要.to(device),否则loss.backward()会成功但梯度为 empty
一个易错细节:torch::nn::Linear 默认初始化权重在 CPU,如果你把 model 移到 CUDA,但忘记调用 model->to(device),它的参数仍在 CPU,而 input 在 GPU,matmul 会失败或静默降级——务必显式迁移整个 module。
C++免费学习笔记(深入):立即使用
在学习笔记中,你将探索 C++ 的入门与实战技巧!











