Pytorch Lighting 的 model.to(device)

Posted 2023-03-28

【中文标题】Pytorch Lighting 的 model.to(device)

【英文标题】：model.to(device) for Pytorch Lighting

【发布时间】：2021-03-18 23:43:29

【问题描述】：

我目前使用 Pytorch Lightning 使用 GPU 训练我的模型

trainer = pl.Trainer( gpus=[0,1],  
        distributed_backend='ddp', 
        resume_from_checkpoint=hparams["resume_from_checkpoint"])
    
    trainer.fit(model, train_dataloader=train_loader, val_dataloaders=val_loader)

关于如何使用定义为使用 GPU 的训练器运行测试样本的说明也很明确

trainer.test(test_dataloader=test_dataloader)

以及如何加载模型并以交互方式使用它

model = transformer.Model.load_from_checkpoint('/checkpoints/run_300_epoch_217.ckpt')
results = model(in_data,

我使用后者通过 docker 容器中的套接字与交互式系统交互。

有没有合适的方法让这个 Pytorch Lightning 模型在 GPU 上运行？ 闪电指令说不要使用 model.to(device)，但它似乎像 Pytorch 一样工作。说明避免副作用的原因？

我开始阅读有关 ONNX 的文章，但我希望有一种简单的方法来指定 GPU，因为交互式设置与 cpu 完美配合。

【参考方案1】：

我的理解是“删除任何 .cuda() 或 to.device() 调用”仅适用于 Lightning 培训师，因为培训师自己处理。

如果您不使用训练器，LightningModule 模块基本上只是一个带有一些命名约定的常规 PyTorch 模型。所以使用model.to(device)是如何在GPU上运行的。