pytorch单机多卡训练

pytorch单机多卡训练

训练

只需要在model定义处增加下面一行：

model = model.to(device)  # device为0号
model = torch.nn.DataParallel(model)

载入模型

  • 如果是多GPU载入，没有问题
  • 如果训练时是多GPU，但是测试时是单GPU，会出现报错 解决办法
posted @ 2020-06-24 15:18  happy_fan  阅读(2214)  评论(0)  编辑  收藏  举报