Skip to content

Commit afbe61a

Browse files
LebranceBWzhulin.zzz
andauthored
doc(pytorch): update docs/pytorch.md (#649)
Co-authored-by: zhulin.zzz <[email protected]>
1 parent f2183a9 commit afbe61a

File tree

1 file changed

+35
-0
lines changed

1 file changed

+35
-0
lines changed

docs/pytorch.md

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,41 @@ torch.Size([1, 2, 28]) # 参数代表在哪里添加维度 0
198198
torch.Size([2, 28, 1])
199199
```
200200

201+
Cuda 相关
202+
---
203+
### 检查 Cuda 是否可用
204+
```python
205+
>>> import torch.cuda
206+
>>> torch.cuda.is_available()
207+
>>> True
208+
```
209+
### 列出 GPU 设备
210+
```python
211+
import torch
212+
device_count = torch.cuda.device_count()
213+
print("CUDA 设备")
214+
for i in range(device_count):
215+
device_name = torch.cuda.get_device_name(i)
216+
total_memory = torch.cuda.get_device_properties(i).total_memory / (1024 ** 3)
217+
print(f"├── 设备 {i}: {device_name}, 容量: {total_memory:.2f} GiB")
218+
print("└── (结束)")
219+
```
220+
### 将模型、张量等数据在 GPU 和内存之间进行搬运
221+
```python
222+
import torch
223+
# Replace 0 to your GPU device index. or use "cuda" directly.
224+
device = f"cuda:0"
225+
# Move to GPU
226+
tensor_m = torch.tensor([1, 2, 3])
227+
tensor_g = tensor_m.to(device)
228+
model_m = torch.nn.Linear(1, 1)
229+
model_g = model_m.to(device)
230+
# Move back.
231+
tensor_m = tensor_g.cpu()
232+
model_m = model_g.cpu()
233+
```
234+
235+
201236
导入 Imports
202237
---
203238

0 commit comments

Comments
 (0)