一、CPU与GPU
- CPU（Central Processing Unit, 中央处理器）：主要包括控制器和运算器
- GPU(Graphics Processing Unit, 图形处理器)：处理统一的，无依赖的大规模数据运算
二、转换数据类型/设备
- tensor.to(*args, **kwargs)，不执行inplace
- module.to(*args, **kwargs)，执行inplace
x = torch.ones((3, 3))
x = x.to(torch.float64)
x = torch.ones((3, 3))
x = x.to("cuda")
linear = nn.Linear(2, 2)
linear.to(torch.double)
gpu1 = torch.device("cuda")
linear.to(gpu1三、 torch.cuda
- torch.cuda.device_count()：计算当前可见可用gpu数
- torch.cuda.get_device_name()：获取gpu名称
- torch.cuda.manual_seed()：为当前gpu设置随机种子
- torch.cuda.manual_seed_all()：为所有可见可用gpu设置随机种子
- 5torch.cuda.set_device()：设置主gpu为哪一个物理gpu（不推荐）
四、并行机制
1、功能：包装模型，实现分发并行机制
torch.nn.DataParallel(module, device_ids=None,output_device=None, dim=0主要参数：
• module: 需要包装分发的模型
• device_ids: 可分发的gpu，默认分发到所有可见可用gpu
• output_device: 结果输出设备
2、查询当前gpu内存剩余
def get_gpu_memory():
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')
return memory_gpu
example:
gpu_memory = get_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)
print("\ngpu free memory: {}".format(gpu_memory))
print("CUDA_VISIBLE_DEVICES :{}".format(os.environ["CUDA_VISIBLE_DEVICES"]))
>>> gpu free memory: [10362, 10058, 9990, 9990]
>>> CUDA_VISIBLE_DEVICES :0,1,3,2