【pytorch使用GPU】

pytorch使用GPU

import torch
import os
import numpy as np

# Pytorch中使用指定的GPU 
os.environ['CUDA_VISIBLE_DEVICE']='3'   # 也可在终端中运行时指定 CUDA_VISIBLE_DEVICES=1 python detect.py

a = torch.empty((2, 3))
print(a)
b = torch.ones((2, 3), device='cuda')
print(b)
pass
a_list = range(16)
c = torch.tensor(a_list).reshape(4, -1)
print(c)

# Tensor 和 Numpy都是矩阵，区别是前者可以在GPU上运行，后者只能在CPU上；
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
d = c.to(device)
print(d)


'''
output:
tensor([[ 4.4162e+36,  4.5616e-41,  5.8214e+37],
        [ 4.5616e-41, -1.3754e-36,  4.5616e-41]])
tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]], device='cuda:1')
'''

print(torch.cuda.is_available())
 
print(torch.cuda.device_count())
 
print(torch.cuda.get_device_name(0))
 
print(torch.cuda.current_device())
 
print(torch.cuda.get_device_properties(0))