Pytorch快速使用手册 01基本配置

import os
import tqdm
import torch
import random
import shutil

import numpy as np

1. Base Config

1.1 Check version of pytorch

print("torch version ",torch.__version__)               # PyTorch version
print("cuda version ", torch.version.cuda)              # Corresponding CUDA version
print("cuDNN version ", torch.backends.cudnn.version())  # Corresponding cuDNN version
print("gpu type ", torch.cuda.get_device_name(0))   # GPU type

torch version 1.12.0+cu116
cuda version 11.6
cuDNN version 8302
gpu type NVIDIA GeForce GTX 1660 Ti

1.2 Update pytroch

pip install --update torch

1.3 Set random number

torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

1.4 Run a tensor in Specified GPU

a = torch.rand(2,2)
a = a.to("cuda:0")
print(a)

tensor([[0.3074, 0.6341],
[0.4901, 0.8964]], device='cuda:0')

1.5 CUDA support ?

print("cuda ok? ", torch.cuda.is_available())

cuda ok? True

1.6 Set mode with cuDNN benchmark

Benchmark will accelarate up train speed, but each there exists random situlation in each train step.

torch.backends.cudnn.benchmark = True

To avoid this situlaiton, you should run the code

torch.backends.cudnn.deterministic = True

1.7 Clear cache of GPU

torch.cuda.empty_cache()

OR

find the progress and kill it

ps aux | grep python

kill -9 [pid]

OR

use nvidia cmd

nvidia-smi --gpu-reset -i [gpu_id]

Github