Pytorch设备选择,多GPU训练

最新推荐文章于 2024-05-09 14:14:03 发布

Always066

最新推荐文章于 2024-05-09 14:14:03 发布

阅读量1.5k

点赞数 1

分类专栏： Pytorch学习笔记

本文链接：https://blog.csdn.net/weixin_43869493/article/details/107139494

版权

Pytorch学习笔记专栏收录该内容

13 篇文章 0 订阅

订阅专栏

To函数

功能：转换数据类型\设备

1.tensor.to()
2.mouble.to()

例子：
在这里插入图片描述

torch.cuda

在这里插入图片描述

多GPU分发并行机制

把数据等分，给不同的GPU运行

torch.nn.DataParallel

在这里插入图片描述

# ============================ 手动选择gpu
gpu_list = [2,3]                                                     
gpu_list_str = ','.join(map(str, gpu_list))                          
os.environ.setdefault("CUDA_VISIBLE_DEVICES", gpu_list_str)          
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ============================ 依内存情况自动选择主gpu
def get_gpu_memory():                                                             
    import platform                                                               
    if 'Windows' != platform.system():                                            
        import os                                                                 
        os.system('nvidia-smi -q -d Memory | grep -A4 GPU | grep Free > tmp.txt') 
        memory_gpu = [int(x.split()[2]) for x in open('tmp.txt', 'r').readlines()]
        os.system('rm tmp.txt')                                                   
    else:                                                                         
        memory_gpu = False                                                        
        print("显存计算功能暂不支持windows操作系统")                                            
    return memory_gpu                                                             
                                                                                  
                                                                                  
gpu_memory = get_gpu_memory()                                                     
if not gpu_memory:                                                                
    print("\ngpu free memory: {}".format(gpu_memory))                             
    gpu_list = np.argsort(gpu_memory)[::-1]                                       
                                                                                  
    gpu_list_str = ','.join(map(str, gpu_list))                                   
    os.environ.setdefault("CUDA_VISIBLE_DEVICES", gpu_list_str)                   
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

常见报错

1.原因：gpu训练出来的model，用map_location映射到cpu上
在这里插入图片描述

path_state_dict = "./model_in_gpu_0.pkl" 
state_dict_load = torch.load(path_state_dict, map_location="cpu")
print("state_dict_load:\n{}".format(state_dict_load))

2.多GPU（DataParallel）训练出来的模型，需要更改一下层的名字，因为在前面多了一个moudle字段
在这里插入图片描述

from collections import OrderedDict                   
new_state_dict = OrderedDict()                        
for k, v in state_dict_load.items():                  
    namekey = k[7:] if k.startswith('module.') else k 
    new_state_dict[namekey] = v                       
print("new_state_dict:\n{}".format(new_state_dict))   
                                                      
net.load_state_dict(new_state_dict)