使用Tesla P4 双卡配置torch机器学习环境

241 阅读2分钟

写在前面

最近一个朋友自己组了一台机器,双Tesla P4 N卡,我准备尝试学习一下配置PyTorch的机器学习环境。主要是之前都是自己的笔记本的游戏显卡,没碰到过这种型号,而且还是双卡,这就更得研究研究环境配置了。

环境

系统:ubuntu 22.04

显卡:Tesla P4 * 2

驱动配置

首先,对于N卡,肯定要先配置N卡驱动。先看一下机器上的显卡

lspci | grep -i nvidia

可以看到是两张Tesla P4卡,这个单卡是8G的显存,根据英伟达的官方数据对比计算能力在6.1 参考:developer.nvidia.com/cuda-gpus

相对来说性价比还是不错的,知名云厂商GPU云服务器对应的价格应该是4k+/月。而我朋友组装的成本价远低于这个月租金,且CPU和内存配置还远超于它。

有点跑题了......继续

不同型号的卡驱动也不同,我们通过这个命令来让ubuntu自己推荐驱动

ubuntu-drivers devices

找到有recommended标记的,就是它了

sudo ubuntu-drivers autoinstall

通过ubuntu-drivers工具自动安装

随后需要reboot重启一下服务器即可

输入nvidia-smi来查看显卡情况

nvidia-smi

至此配套的显卡驱动装好了,记住这个CUDA版本信息。

PyTorch配置

刚刚咱们的CUDA版本信息是11.4,目前官方给出11.x版本最新的是11.8,官网链接:pytorch.org/

经过查询,CUDA是支持向下兼容的

我这个装的驱动版本号是470,所以直接使用11.8是没问题的

参考链接:docs.nvidia.com/cuda/cuda-t…

这就简单了,直接复制该命令安装即可。

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

随后打开python3,执行

import torch
print(torch.cuda.is_available())
print(torch.cuda.device_count())

至此,torch配置完毕

使用双卡运行

这里简单提一嘴,如果想使用多卡同时调用,需要使用DataParallel方法即可,这是单机多卡的适用方法

model = Model()
model = torch.nn.DataParallel(model)
model.to('cuda')

这样就会自动的将模型分配到所有可用的GPU上了。