ai

nvidia-docker运行pytorch-gpu

1.宿主机器安装显卡驱动和cuda 11.8, cudnn 11.8

2.拉镜像(基于ubuntu):

docker pull anibali/pytorch:2.0.1-cuda11.8

运行容器:

docker run -it --init   --gpus=all   --ipc=host --name pytorch -p 1778:8888 --volume="$PWD:/app"   [镜像id] python3

运行一下脚本检查cuda是否可用:

import torch
flag = torch.cuda.is_available()
print(flag)

ngpu= 1
# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda()) 

正确输出结果类似:

True
cuda:0
GeForce GTX 1080
tensor([[0.9530, 0.4746, 0.9819],
        [0.7192, 0.9427, 0.6768],
        [0.8594, 0.9490, 0.6551]], device='cuda:0')


关于作者

程序员,软件工程师,java, golang, rust, c, python,vue, Springboot, mybatis, mysql,elasticsearch, docker, maven, gcc, linux, ubuntu, centos, axum,llm, paddlepaddle, onlyoffice,minio,银河麒麟,中科方德,rpm