• 首页 首页 icon
  • 工具库 工具库 icon
    • IP查询 IP查询 icon
  • 内容库 内容库 icon
    • 快讯库 快讯库 icon
    • 精品库 精品库 icon
    • 问答库 问答库 icon
  • 更多 更多 icon
    • 服务条款 服务条款 icon

模型加载至 cpu 和 gpu 的方式

武飞扬头像
M_TDM
帮助1

  1. 采用 from_pretrained 的方式,模型正常情况下,BertMoldel.from_pretrained() 是会 load 在 cpu 上的,内部 map_location 默认设置成 cpu,如果想要部署在gpu,执行下面三句话。

  1.  
    BertMoldel.from_pretrained()
  2.  
    device=torch.device(’cuda’)
  3.  
    model.to(device)
  1. 采用 load_state_dict 的方式加载模型,模型是部署在 哪里可以指定,如果想部署到 gpu,无需修改第一行,直接再加入4.5行。

  1.  
    state_dict=torch.load(model_path, map_location=’cpu’)
  2.  
    #部署到 gpu,把上面改为map_location=’gpu’
  3.  
    model.load_state_dict(state_dict)
  4.  
    #已在CPU上加载,下面两句也可加入GPU
  5.  
    device=torch.device(’cuda’)
  6.  
    model.to(device)

这篇好文章是转载于:学新通技术网

  • 版权申明: 本站部分内容来自互联网,仅供学习及演示用,请勿用于商业和其他非法用途。如果侵犯了您的权益请与我们联系,请提供相关证据及您的身份证明,我们将在收到邮件后48小时内删除。
  • 本站站名: 学新通技术网
  • 本文地址: /boutique/detail/tanhghbbbi
系列文章
更多 icon
同类精品
更多 icon
继续加载