モデルが CPU と GPU にロードされる方法

  1. from_pretrained メソッドを使用します。通常の状況では、BertMoldel.from_pretrained() が CPU にロードされます。内部の map_location はデフォルトで CPU に設定されます。GPU にデプロイしたい場合は、次の 3 つの文を実行します。

BertMoldel.from_pretrained()
device=torch.device(’cuda’)
model.to(device) 
  1. モデルをロードするには、load_state_dict メソッドを使用します。モデルがデプロイされる場所を指定できます。GPU にデプロイする場合は、最初の行を変更する必要はなく、行 4.5 を追加するだけです。

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

おすすめ

転載: blog.csdn.net/M_TDM/article/details/129436122