Pytorch2.0中compiled_model=torch.compile(model) 的正确添加位置

黄渡猿

已于 2023-03-18 12:34:41 修改

阅读量3.5k

点赞数 2

文章标签：深度学习 pytorch 人工智能计算机视觉 python

于 2023-03-18 12:30:30 首次发布

本文链接：https://blog.csdn.net/CarryEKAIruiui/article/details/129635139

版权

PyTorch官网推出了2.0稳定版，新特性torch.compile允许用户只需添加model=torch.compile(model)即可实现模型训练加速。无论是CPU训练、单卡训练还是多卡训练(DistributedDataParallel)，只需关注这行代码的位置，无需大量修改原有代码，即可享受性能提升。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

今天pytorch官网更新了pytorch2.0稳定版，迫不及待的我直接更新了，确实像官方所说，只需加入model=torch.compile(model)一行代码即可加速，加入的位置如下。

cpu训练：

 model=UNet(deep_supervision=True)
 model=torch.compile(model)

单卡训练：

 model=UNet(deep_supervision=True)
 model.to(Device)
 model=torch.compile(model)

多卡训练：

 model=UNet(deep_supervision=True)
 model.to(Device)
 model=nn.parallel.DistributedDataParallel(
         model,
         device_ids=[local_rank],
         output_device=local_rank,
         broadcast_buffers=False,
     )
 model=torch.compile(model)

注意 model = torch.compile(model) 这句话的位置对了就可以了，其他的不用改！！