model.to（设备）用于 Pytorch 照明

Question

我目前使用 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 容器中的套接字与交互式系统交互。

有没有合适的方法在 GPU 上制作这个 Pytorch Lightning 模型运行？ Lightning 说明说不要使用 model.to(device)，但它看起来就像 Pytorch 一样工作。说明避免副作用的原因？

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

Answer 1

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

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

model.to（设备）用于 Pytorch 照明

model.to(device) for Pytorch Lighting

pytorch

pytorch-lightning