PyTorch中的CUDA的操作方法

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

print(torch.__version__)
1.12.1+cu116

print(torch.cuda.is_available())
True

print("default device: {}".format(torch.Tensor([4,5,6]).device))
default device: cpu

print("available cpu devices: {}".format(torch.cuda.os.cpu_count()))
available cpu devices: 20

print("available gpu devices: {}".format(torch.cuda.device_count()))
available gpu devices: 1

print("gpu device name: {}".format(torch.cuda.get_device_name(torch.device("cuda:0"))))
gpu device name: NVIDIA GeForce GTX 1080 Ti

device = torch.Tensor([1,2,3], device="cpu:0").device
print("device type: {}".format(device))
device type: cpu

cpu1 = torch.device("cpu:0")
print("cpu device: {}:{}".format(cpu1.type, cpu1.index))
cpu device: cpu:0

gpu = torch.device(0)
print("gpu device: {}:{}".format(gpu.type, gpu.index))
gpu device: cuda:0

gpu = torch.device("cuda:0")
print("gpu device: {}:{}".format(gpu.type, gpu.index))
gpu device: cuda:0

# 默认创建的tensor是在cpu上创建的
cpu_tensor = torch.Tensor([[1,4,7],[3,6,9],[2,5,8]])
print(cpu_tensor.device)

# 通过to方法将cpu_tensor拷贝到gpu上
gpu_tensor1 = cpu_tensor.to(torch.device("cuda:0"))
print(gpu_tensor1.device)

# 通过cuda方法将cpu_tensor拷贝到gpu上
gpu_tensor2 = cpu_tensor.cuda(torch.device("cuda:0"))
print(gpu_tensor2.device)

# 将gpu_tensor2拷贝到cpu上
gpu_tensor3 = cpu_tensor.copy_(gpu_tensor2)
print(gpu_tensor3.device)
print(gpu_tensor3)

def copy_(self, src, non_blocking=False):
    ......
    return _te.Tensor(*(), **{})

gpu_tensor1 = torch.tensor([[2,5,8],[1,4,7],[3,6,9]], device=torch.device("cuda:0"))
print(gpu_tensor1.device)

# 在gpu设备上创建随机数tensor
print(torch.rand((3,4), device=torch.device("cuda:0")))

# 在gpu设备上创建0值tensor
print(torch.zeros((2,5), device=torch.device("cuda:0")))

cuda = torch.device("cuda")
# 创建默认的stream，A就是使用的默认stream
s = torch.cuda.Stream()
A = torch.randn((1,10), device=cuda)
for i in range(100):
    # 在新的stream上对默认的stream上创建的tensor进行求和
    with torch.cuda.stream(s):
        # 存在的问题是：torch.sum()可能会在torch.randn()之前执行
        B = torch.sum(A)
        print(B)

cuda = torch.device("cuda")
s = torch.cuda.Stream()
A = torch.randn((1,10), device=cuda)
default_stream = torch.cuda.current_stream()
print("Default Stream: {}".format(default_stream))
# 等待创建A的stream执行完毕
torch.cuda.Stream.synchronize(default_stream)
for i in range(100):
    # 在新的stream上对默认的stream上创建的tensor进行求和
    with torch.cuda.stream(s):
        print("current stream: {}".format(torch.cuda.current_stream()))
        B = torch.sum(A)
        print(B)

from torch.utils.data._utils.pin_memory import pin_memory
x = torch.Tensor([[1,2,4], [5, 7, 9], [3, 7, 10]])
# 通过pin_memory()方法将x复制到固定缓冲区
y = pin_memory(x)
# 在tensor上直接调用pin_memory()方法将tensor复制到固定缓冲区
z = x.pin_memory()
# id()方法返回tensor的内存地址，pin_memory()返回tensor对象的拷贝，因此内存地址是不同的
print("id: {}".format(id(x)))
print("id: {}".format(id(y)))
print("id: {}".format(id(z)))
# 当tensor放入固定缓冲区后，就可以异步将数据复制到gpu设备上了
a = z.cuda(non_blocking=True)
print(a)
print("is_pinned: {}/{}".format(x.is_pinned(), z.is_pinned()))

device = torch.device("cpu")
if torch.cuda.is_available():
    device = torch.device("cuda")
a = torch.tensor([1,2,3], device=device)
print(a)

class LinearRegression(torch.nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.linear = torch.nn.Linear(1, 1)
    def forward(self, x):
        return self.linear(x)
regression = LinearRegression().to(device=device)
for param in regression.parameters():
    print(param)