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

# 默認(rèn)創(chuàng)建的tensor是在cpu上創(chuàng)建的
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設(shè)備上創(chuàng)建隨機(jī)數(shù)tensor
print(torch.rand((3,4), device=torch.device("cuda:0")))

# 在gpu設(shè)備上創(chuàng)建0值tensor
print(torch.zeros((2,5), device=torch.device("cuda:0")))

cuda = torch.device("cuda")
# 創(chuàng)建默認(rèn)的stream，A就是使用的默認(rèn)stream
s = torch.cuda.Stream()
A = torch.randn((1,10), device=cuda)
for i in range(100):
    # 在新的stream上對默認(rèn)的stream上創(chuàng)建的tensor進(jìn)行求和
    with torch.cuda.stream(s):
        # 存在的問題是：torch.sum()可能會在torch.randn()之前執(zhí)行
        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))
# 等待創(chuàng)建A的stream執(zhí)行完畢
torch.cuda.Stream.synchronize(default_stream)
for i in range(100):
    # 在新的stream上對默認(rèn)的stream上創(chuàng)建的tensor進(jìn)行求和
    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復(fù)制到固定緩沖區(qū)
y = pin_memory(x)
# 在tensor上直接調(diào)用pin_memory()方法將tensor復(fù)制到固定緩沖區(qū)
z = x.pin_memory()
# id()方法返回tensor的內(nèi)存地址，pin_memory()返回tensor對象的拷貝，因此內(nèi)存地址是不同的
print("id: {}".format(id(x)))
print("id: {}".format(id(y)))
print("id: {}".format(id(z)))
# 當(dāng)tensor放入固定緩沖區(qū)后，就可以異步將數(shù)據(jù)復(fù)制到gpu設(shè)備上了
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)