GPU计算¶
In [1]:
Copied!
# 查看显卡信息
!nvidia-smi
# 查看显卡信息
!nvidia-smi
Thu Jan 29 15:04:13 2026
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 566.24 Driver Version: 566.24 CUDA Version: 12.7 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Driver-Model | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 4060 ... WDDM | 00000000:01:00.0 Off | N/A |
| N/A 47C P0 11W / 80W | 0MiB / 8188MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| No running processes found |
+-----------------------------------------------------------------------------------------+
计算设备¶
In [5]:
Copied!
import torch
from torch import nn
print(torch.cuda.is_available()) # 是否能使用cuda
print(torch.cuda.device_count()) # 查看GPU数量
print(torch.cuda.current_device()) # 查看当前GPU的索引号
print(torch.cuda.get_device_name(0)) # 根据索引查看当前GPU的名字
import torch
from torch import nn
print(torch.cuda.is_available()) # 是否能使用cuda
print(torch.cuda.device_count()) # 查看GPU数量
print(torch.cuda.current_device()) # 查看当前GPU的索引号
print(torch.cuda.get_device_name(0)) # 根据索引查看当前GPU的名字
True 1 0 NVIDIA GeForce RTX 4060 Laptop GPU
Tensor的GPU计算¶
In [6]:
Copied!
x = torch.tensor([1, 2, 3])
x = x.cuda() # .cuda()将Tensor复制到GPU上
x
x = torch.tensor([1, 2, 3])
x = x.cuda() # .cuda()将Tensor复制到GPU上
x
Out[6]:
tensor([1, 2, 3], device='cuda:0')
In [7]:
Copied!
x.device # 查看Tensor所在的设备
x.device # 查看Tensor所在的设备
Out[7]:
device(type='cuda', index=0)
In [8]:
Copied!
# 在创建时指定GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 以下两种方法二选一
x = torch.tensor([1, 2, 3], device=device)
x = torch.tensor([1, 2, 3]).to(device)
x
# GPU上的tensor之间运算，结果还是在GPU上
# 但GPU的tensor不能和CPU的tensor直接运算
# 在创建时指定GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 以下两种方法二选一
x = torch.tensor([1, 2, 3], device=device)
x = torch.tensor([1, 2, 3]).to(device)
x
# GPU上的tensor之间运算，结果还是在GPU上
# 但GPU的tensor不能和CPU的tensor直接运算
Out[8]:
tensor([1, 2, 3], device='cuda:0')
模型的GPU计算¶
In [9]:
Copied!
net = nn.Linear(3, 1)
net.cuda() # 同样用.cuda()复制到GPU上
list(net.parameters())[0].device
# 需要保证模型输入的Tensor和模型都在同一设备上，否则会报错
net = nn.Linear(3, 1)
net.cuda() # 同样用.cuda()复制到GPU上
list(net.parameters())[0].device
# 需要保证模型输入的Tensor和模型都在同一设备上，否则会报错
Out[9]:
device(type='cuda', index=0)