OpenFOAM libtorch tutorial step by step

xyl308120

老哥，我把东岳老师那个1D槽道流的代码用AI改成了CUDA版本。然后在https://download.pytorch.org/libtorch/cu117 （我的cuda版本是11.7）下载的gpu版本的libtorch。现在发现已经可以用gpu训练了，但是训练速度还不如cpu，慢不少。准备后续再捣鼓一下。我机器配置 GPU： nvidia T400 4G， CPU：I7-12700，供参考

#include <torch/torch.h>
#include <torch/cuda.h> // 添加CUDA支持

using namespace std;

class NN : public torch::nn::Module {
    torch::nn::Sequential net_;
public:
    NN() {
        net_ = register_module("net", torch::nn::Sequential(
            torch::nn::Linear(1,8),
            torch::nn::Tanh(),
            torch::nn::Linear(8,1)
        ));
    }
    torch::Tensor forward(torch::Tensor x) {
        return net_->forward(x);
    }
};

int main() {
    // 设备检测与设置
    torch::Device device = torch::kCPU;
    if (torch::cuda::is_available()) {
        std::cout << "CUDA可用，使用GPU加速训练" << std::endl;
        device = torch::kCUDA;
    }

    int iter = 1;
    int IterationNum = 1000000;
    double learning_rate = 0.001;
    double dy = 0.00125;
    int cellN = static_cast<int>(0.05/dy);

    // 模型迁移至GPU
    auto model = std::make_shared<NN>();
    model->to(device);
    
    // 优化器
    auto opti = std::make_shared<torch::optim::AdamW>(
        model->parameters(), 
        torch::optim::AdamWOptions(learning_rate)
    );

    // 数据迁移至GPU (关键步骤！)
    auto init = torch::full({}, 1.0).to(device);
    auto mesh = torch::arange(-0.05, 0, dy, torch::requires_grad())
                .unsqueeze(1).to(device);
    auto dpdxByNu = torch::ones({cellN, 1}, torch::kFloat).to(device);
    auto nut = torch::zeros({cellN}).to(device);
        
    for (int i = 0; i < IterationNum; i++) {
        opti->zero_grad();
        auto upred = model->forward(mesh);
        
        // 梯度计算（确保权重张量在GPU）
        auto dudy = torch::autograd::grad(
            {upred}, {mesh}, 
            {torch::ones_like(upred).to(device)}, true, true
        )[0];
        
        auto dudyy = torch::autograd::grad(
            {dudy}, {mesh}, 
            {torch::ones_like(upred).to(device)}, true, true
        )[0];

        auto dudyTop = dudy[cellN - 1][0];
        auto meanU = torch::mean(upred);
        auto ubottom = torch::full_like(upred[0], 0.0375).to(device); // 边界条件迁移
        auto dpdxByNu = torch::full({cellN}, 1.2e-05/1e-08).to(device);
        
        // 损失计算
        auto loss_bottom = torch::mse_loss(upred[0], ubottom);
        auto loss_top = torch::mse_loss(dudyTop, torch::zeros_like(dudyTop));
        auto loss_bd = loss_bottom + 1000*loss_top;
        auto loss_ini = 1000*torch::mse_loss(meanU, init);
        auto loss_pde = torch::mse_loss(-dudyy.reshape({cellN}), dpdxByNu);
        auto loss = loss_pde + loss_bd + loss_ini;

        // 反向传播
        loss.backward();
        opti->step();
        
        // 输出需同步到CPU
        if (iter % 500 == 0) {
            double loss_value = loss.item<double>();
            cout << iter << " loss = " << loss_value << endl;
        }
        iter++;
    }

    torch::save(model, "model.pth");
    std::cout << "训练完成！" << std::endl;
    return 0;
}

李东岳

现在发现已经可以用gpu训练了，但是训练速度还不如cpu，慢不少。

看你GPU占用多少。

xyl308120

@李东岳
东岳老师你好，我训练的时候GPU占比%38左右。我还想顺便请教一下

loss_ini = 1000*crit(meanU, init)

这个损失的意义。因为我发现1D槽道流openfoam基准数据的平均值差不多就是1左右，而这里相当于是要对训练添加一个整体约束之类的东西吗。

李东岳

是的，一个整体平均值约束。

增加神经元看看能不能把GPU拉上100%

CFD中文网

OpenFOAM libtorch tutorial step by step