Con el método from_pretrained, BertMoldel.from_pretrained() se cargará en la CPU en condiciones normales del modelo, y map_location interno se establece en la CPU de forma predeterminada. Si desea implementarlo en la GPU, ejecute las siguientes tres oraciones.
BertMoldel.from_pretrained()
device=torch.device(’cuda’)
model.to(device) Use el método load_state_dict para cargar el modelo. Puede especificar dónde se implementa el modelo. Si desea implementar en la GPU, no necesita modificar la primera línea, solo agregue la línea 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)