map_location参数在torch.load函数中用于指定加载模型时的设备映射。它的作用是控制模型权重被加载到哪个设备上。在PyTorch中,设备可以是CPU或GPU,而GPU可以有多个,每个都有其独立的索引。map_location参数的灵活使用能够让我们轻松地在不同设备之间迁移模型,从而充分利用不同设备的计算优势。
CPU加载:当你想在CPU上加载模型时,可以设置map_location="cpu"。这适用于那些不需要GPU加速的推理任务,或者在没有GPU的环境中部署模型。
指定GPU加载:如果你有多个GPU,并且想将模型加载到特定的GPU上,可以使用"cuda:X"格式的字符串,其中X是GPU的索引。这在多GPU环境中非常有用,可以确保模型加载到指定的设备上。
自动选择GPU:如果你只想在GPU上加载模型,但不关心具体是哪一个GPU,可以设置map_location=torch.device("cuda")。这会自动选择第一个可用的GPU来加载模型。
{"cuda:1":"cuda:0"}:将原本在GPU 1上的张量加载到GPU 0上,例如:
# map_location参数为一个字典,将原本在GPU 1上的张量映射到GPU 0上
tensors = torch.load("tensors.pt", map_location={"cuda:1":"cuda:0"})