当前位置：首页 > article >正文

《PyTorch基础教程》01 搭建环境基于Docker搭建ubuntu22+Python3.10+Pytorch2+cuda11+jupyter的开发环境

article 2025/3/19 6:18:49

01 环境搭建

Docker部署PyTorch

拉取cnstark/pytorch镜像

拉取镜像：

docker pull cnstark/pytorch:2.0.1-py3.10.11-cuda11.8.0-ubuntu22.04

导出镜像：

docker save -o pytorch2_python310_cuda11_ubuntu22.tar cnstark/pytorch:2.0.1-py3.10.11-cuda11.8.0-ubuntu22.04

导入镜像：

docker load -i pytorch2_python310_ubuntu22.tar

运行镜像：

mkdir -p /docker/pytorch/project
mkdir -p /docker/pytorch/dataset
docker run --name pytorch -itd -v /docker/pytorch/project:/workspace -v /docker/pytorch/dataset:/workspace/dataset cnstark/pytorch:2.0.1-py3.10.11-cuda11.8.0-ubuntu22.04

# 开启GPU
docker run --name pytorch --gpus all -itd -v /docker/pytorch/project:/workspace -v /docker/pytorch/dataset:/workspace/dataset cnstark/pytorch:2.0.1-py3.10.11-cuda11.8.0-ubuntu22.04

测试PyTorch脚本

main.py

# -*- coding: utf-8 -*-

import torch

dtype = torch.FloatTensor
# dtype = torch.cuda.FloatTensor # Uncomment this to run on GPU

# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 1000, 100, 10

# Create random input and output data
x = torch.randn(N, D_in).type(dtype)
y = torch.randn(N, D_out).type(dtype)

# Randomly initialize weights
w1 = torch.randn(D_in, H).type(dtype)
w2 = torch.randn(H, D_out).