Google Colab TPU 使用方法

最新推荐文章于 2025-04-04 18:09:29 发布

__init__:

最新推荐文章于 2025-04-04 18:09:29 发布

阅读量358

点赞数 4

文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/qq_45802280/article/details/144251689

版权

最近在白嫖Google colab，发现有免费的TPU可以。使用方式也很简单：

首先需要安装TPU支持的库：

 # 安装最新的 torch_xla 库
 pip install torch_xla -f <https://storage.googleapis.com/tpu-pytorch/wheels/colab.html>

随后只需要指定计算device即可愉快的玩耍啦：

 import torch
 import torch_xla
 import torch_xla.core.xla_model as xm
 
 # 设定设备为 TPU
 device = xm.xla_device()
 
 # 创建张量并移动到 TPU 上
 x = torch.randn(3, 3).to(device)
 y = torch.randn(3, 3).to(device)
 
 # 执行张量操作
 z = x + y
 
 # 将结果从 TPU 移回 CPU
 z_cpu = z.cpu()
 print(z_cpu)