pytorch - Pytorch 照明的 model.to(device)

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 容器中的套接字与交互式系统进行交互。

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

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

Answer 1

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

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