一個簡單的矩陣乘法例子來演示在 PyTorch 中如何針對 GPU 和 TPU 使用不同的處理方式。

這個例子會展示核心的區別在于如何獲取和指定計算設備，以及（對于 TPU）可能需要額外的庫和同步操作。

示例代碼：

import torch
import time# --- GPU 示例 ---
print("--- GPU 示例 ---")
# 檢查是否有可用的 GPU (CUDA)
if torch.cuda.is_available():gpu_device = torch.device('cuda')print(f"檢測到 GPU。使用設備: {gpu_device}")# 創建張量并移動到 GPU# 在張量創建時直接指定 device='cuda' 或 .to('cuda')tensor_a_gpu = torch.randn(1000, 2000, device=gpu_device)tensor_b_gpu = torch.randn(2000, 1500, device=gpu_device)# 在 GPU 上執行矩陣乘法start_time = time.time()result_gpu = torch.mm(tensor_a_gpu, tensor_b_gpu)torch.cuda.synchronize() # 等待 GPU 計算完成end_time = time.time()print(f"在 GPU 上執行了矩陣乘法，結果張量大小: {result_gpu.shape}")print(f"GPU 計算耗時: {end_time - start_time:.4f} 秒")# print(result_gpu) # 可以打印結果，但對于大張量會很多else:print("未檢測到 GPU。無法運行 GPU 示例。")# --- TPU 示例 ---
print("\n--- TPU 示例 ---")
# 導入 PyTorch/XLA 庫
# 注意：這個庫需要在支持 TPU 的環境 (如 Google Colab TPU runtime 或 Cloud TPU VM) 中安裝和運行
try:import torch_xlaimport torch_xla.core.xla_model as xmimport torch_xla.distributed.parallel_loader as plimport torch_xla.distributed.xla_multiprocessing as xmp# 檢查是否在 XLA (TPU) 環境中if xm.xla_device() is not None:IS_TPU_AVAILABLE = Trueelse:IS_TPU_AVAILABLE = Falseexcept ImportError:print("未找到 torch_xla 庫。")IS_TPU_AVAILABLE = False
except Exception as e:print(f"初始化 torch_xla 失敗: {e}")IS_TPU_AVAILABLE = Falseif IS_TPU_AVAILABLE:# 獲取 TPU 設備tpu_device = xm.xla_device()print(f"檢測到 TPU。使用設備: {tpu_device}")# 創建張量并移動到 TPU (通過 XLA 設備)# 在張量創建時直接指定 device=tpu_device 或 .to(tpu_device)# 注意：TPU 操作通常是惰性的，數據和計算可能會在 xm.mark_step() 或其他同步點時才實際執行tensor_a_tpu = torch.randn(1000, 2000, device=tpu_device)tensor_b_tpu = torch.randn(2000, 1500, device=tpu_device)# 在 TPU 上執行矩陣乘法 (通過 XLA)start_time = time.time()result_tpu = torch.mm(tensor_a_tpu, tensor_b_tpu)# 觸發執行和同步 (TPU 操作通常是惰性的，需要顯式步驟來編譯和執行)# 在實際訓練循環中，通常在一個 minibatch 結束時調用 xm.mark_step()xm.mark_step()# 注意：TPU 的時間測量可能需要通過特定 XLA 函數，這里使用簡單的 time() 可能不精確反映 TPU 計算時間end_time = time.time()print(f"在 TPU 上執行了矩陣乘法，結果張量大小: {result_tpu.shape}")#print(f"TPU (包含編譯和同步) 耗時: {end_time - start_time:.4f} 秒") # 這里的計時僅供參考# print(result_tpu) # 可以打印結果else:print("無法運行 TPU 示例，因為未找到 torch_xla 庫 或 不在 TPU 環境中。")print("要在 Google Colab 中運行 TPU 示例，請在 'Runtime' -> 'Change runtime type' 中選擇 TPU。")

代碼解釋：

1. 導入: 除了 torch，GPU 示例不需要額外的庫。但 TPU 示例需要導入 torch_xla 庫。
2. 設備獲取:
   • GPU 使用 torch.device('cuda') 或更簡單的 'cuda' 字符串來指定設備。torch.cuda.is_available() 用于檢查 CUDA 是否可用。
   • TPU 使用 torch_xla.core.xla_model.xla_device() 來獲取 XLA 設備對象。通常需要檢查 torch_xla 是否成功導入以及 xm.xla_device() 是否返回一個非 None 的設備對象來確定 TPU 環境是否可用。
3. 張量創建/移動:
   • 無論是 GPU 還是 TPU，都可以通過在創建張量時指定 device=... 或使用 .to(device) 方法將已有的張量移動到目標設備上。
4. 計算: 執行矩陣乘法 torch.mm() 的代碼在兩個例子中看起來是相同的。這是 PyTorch 的一個優點，上層代碼在不同設備上可以保持相似。
5. 同步:
   • GPU 操作在調用時通常是異步的，但 torch.cuda.synchronize() 會阻塞 CPU，直到所有 GPU 操作完成，這在計時時是必需的。
   • TPU 操作通過 XLA 編譯和執行，通常是惰性的 (lazy)。這意味著調用 torch.mm() 可能只是構建計算圖，實際計算可能不會立即發生。xm.mark_step() 是一個重要的同步點，它會觸發 XLA 編譯當前構建的計算圖并在 TPU 上執行，然后等待執行完成。在實際訓練循環中，這通常在每個 mini-batch 結束時調用。

核心區別在于設備層面的處理方式： 原生 PyTorch 直接通過 CUDA API 與 GPU 交互，而對 TPU 的支持則需要借助 torch_xla 庫作為中介，通過 XLA 編譯器來生成和管理 TPU 上的執行。