libtorch技巧
0
- 使用
float16计算配置optims.eps(1e-5);可以避免梯度计算nan - 使用
amp混合精度可以减小内存使用 - 推理是建议使用
torch::NoGradGuard避免梯度计算导致内存暴涨 - 显卡内存不够可以加载数据集到内存每次复制到显存里面计算
自动混合精度问题
// 自动混合精度：没有GradScaler容易导致梯度不能下降
this->optimizer->zero_grad();
at::autocast::set_autocast_dtype(torch::kCUDA, torch::kFloat16);
at::autocast::set_autocast_enabled(torch::kCUDA, true);
at::autocast::set_autocast_cache_enabled(true);
auto denoise = this->model->forward(noise_images, steps);
at::autocast::clear_cache();
at::autocast::set_autocast_enabled(torch::kCUDA, false);
loss.backward();
torch::nn::utils::clip_grad_norm_(this->model->parameters(), this->params.grad_clip);
this->optimizer->step();