#!/usr/bin/env python3 """ 使用AMP的TPU训练脚本 正确处理混合精度训练，避免dtype不匹配问题 """ import os import time import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms # 设置AMP相关的环境变量 os.environ['XLA_FLAGS'] = ( '--xla_cpu_multi_thread_eigen=true ' '--xla_cpu_enable_fast_math=true' ) os.environ['XLA_USE_BF16'] = '1' # 启用bf16 import torch_xla.core.xla_model as xm import torch_xla.distributed.parallel_loader as pl import torch_xla.amp as xla_amp class AMPModel(nn.Module): """支持AMP的简单模型""" def __init__(self, input_size=784, hidden_size=512, num_classes=10): super(AMPModel, self).__init__() self.network = nn.Sequential( nn.Linear(input_size, hidden_size), nn.ReLU(inplace=True), nn.Dropout(0.2), nn.Linear(hidden_size, hidden_size // 2), nn.ReLU(inplace=True), nn.Dropout(0.2), nn.Linear(hidden_size // 2, num_classes) ) def forward(self, x): # 展平输入 x = x.view(x.size(0), -1) return self.network(x) class AMPTrainer: """AMP训练器""" def __init__(self, model, device, learning_rate=0.001): self.model = model self.device = device self.optimizer = optim.Adam(model.parameters(), lr=learning_rate) self.criterion = nn.CrossEntropyLoss() # 初始化AMP scaler self.scaler = xla_amp.GradScaler() print(f"✅ AMP训练器初始化完成") print(f" 设备: {device}") print(f" 模型参数: {sum(p.numel() for p in model.parameters()):,}") def train_step(self, data, target): """单个AMP训练步骤""" self.model.train() self.optimizer.zero_grad() # 使用autocast进行混合精度前向传播 with xla_amp.autocast(): output = self.model(data) loss = self.criterion(output, target) # 使用scaler进行反向传播 self.scaler.scale(loss).backward() # 梯度裁剪（可选） self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) # 更新参数 self.scaler.step(self.optimizer) self.scaler.update() # 计算准确率 pred = output.argmax(dim=1) correct = pred.eq(target).sum().item() accuracy = correct / target.size(0) return loss.item(), accuracy def evaluate_step(self, data, target): """单个评估步骤""" self.model.eval() with torch.no_grad(): with xla_amp.autocast(): output = self.model(data) loss = self.criterion(output, target) pred = output.argmax(dim=1) correct = pred.eq(target).sum().item() accuracy = correct / target.size(0) return loss.item(), accuracy def get_mnist_loaders(batch_size=64): """获取MNIST数据加载器""" transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = torchvision.datasets.MNIST( root='./mnist_data', train=True, download=True, transform=transform ) test_dataset = torchvision.datasets.MNIST( root='./mnist_data', train=False, download=True, transform=transform ) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=0 ) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=batch_size, shuffle=False, num_workers=0 ) return train_loader, test_loader def train_with_amp(): """使用AMP进行训练""" print("🚀 开始AMP TPU训练...") # 获取设备 device = xm.xla_device() print(f"📱 设备: {device}") # 创建模型 model = AMPModel(input_size=784, hidden_size=512, num_classes=10).to(device) # 创建训练器 trainer = AMPTrainer(model, device, learning_rate=0.001) # 获取数据 print("📥 加载MNIST数据...") train_loader, test_loader = get_mnist_loaders(batch_size=64) # 使用XLA并行加载器 train_device_loader = pl.MpDeviceLoader(train_loader, device) test_device_loader = pl.MpDeviceLoader(test_loader, device) print("🎯 开始AMP训练...") # 训练循环 num_epochs = 2 train_losses = [] train_accuracies = [] for epoch in range(num_epochs): print(f"\n📊 Epoch {epoch + 1}/{num_epochs}") epoch_start = time.time() epoch_loss = 0.0 epoch_acc = 0.0 num_batches = 0 max_batches_per_epoch = 200 # 限制每个epoch的批次数 for batch_idx, (data, target) in enumerate(train_device_loader): if batch_idx >= max_batches_per_epoch: break # 训练步骤 loss, accuracy = trainer.train_step(data, target) epoch_loss += loss epoch_acc += accuracy num_batches += 1 # 每20个批次同步一次 if batch_idx % 20 == 0: xm.mark_step() avg_loss = epoch_loss / num_batches avg_acc = epoch_acc / num_batches * 100 print(f" 批次 {batch_idx:3d}/{max_batches_per_epoch} | " f"损失: {avg_loss:.4f} | " f"准确率: {avg_acc:.2f}%") # Epoch结束同步 xm.mark_step() xm.wait_device_ops() epoch_time = time.time() - epoch_start final_loss = epoch_loss / num_batches final_acc = epoch_acc / num_batches * 100 train_losses.append(final_loss) train_accuracies.append(final_acc) print(f"✅ Epoch {epoch + 1} 完成 | " f"耗时: {epoch_time:.2f}s | " f"平均损失: {final_loss:.4f} | " f"平均准确率: {final_acc:.2f}%") return trainer, train_losses, train_accuracies def test_with_amp(trainer): """使用AMP进行测试""" print("\n🧪 开始AMP测试...") device = xm.xla_device() _, test_loader = get_mnist_loaders(batch_size=64) test_device_loader = pl.MpDeviceLoader(test_loader, device) total_loss = 0.0 total_acc = 0.0 num_batches = 0 max_test_batches = 100 start_time = time.time() for batch_idx, (data, target) in enumerate(test_device_loader): if batch_idx >= max_test_batches: break loss, accuracy = trainer.evaluate_step(data, target) total_loss += loss total_acc += accuracy num_batches += 1 if batch_idx % 20 == 0: xm.mark_step() xm.mark_step() xm.wait_device_ops() test_time = time.time() - start_time avg_loss = total_loss / num_batches avg_acc = total_acc / num_batches * 100 print(f"✅ 测试完成!") print(f"⏱️ 测试时间: {test_time:.2f}秒") print(f"🎯 测试损失: {avg_loss:.4f}") print(f"🎯 测试准确率: {avg_acc:.2f}%") return avg_loss, avg_acc def main(): """主函数""" print("=" * 60) print("⚡ AMP TPU训练示例") print("=" * 60) try: # 训练 trainer, train_losses, train_accuracies = train_with_amp() # 测试 test_loss, test_acc = test_with_amp(trainer) # 保存模型 print("\n💾 保存模型...") model_cpu = trainer.model.cpu() torch.save({ 'model_state_dict': model_cpu.state_dict(), 'train_losses': train_losses, 'train_accuracies': train_accuracies, 'test_loss': test_loss, 'test_accuracy': test_acc }, 'amp_mnist_model.pth') print("✅ 模型已保存到 amp_mnist_model.pth") print("\n🎉 AMP训练完成!") print(f"📊 最终训练准确率: {train_accuracies[-1]:.2f}%") print(f"📊 测试准确率: {test_acc:.2f}%") if train_accuracies[-1] > 85 and test_acc > 80: print("✅ AMP训练成功! 模型性能优秀") else: print("⚠️ 模型性能一般，但AMP功能正常") except Exception as e: print(f"❌ AMP训练失败: {e}") import traceback traceback.print_exc() print("\n💡 故障排除建议:") print(" 1. 确保PyTorch XLA版本支持AMP") print(" 2. 检查TPU资源是否充足") print(" 3. 尝试减小batch_size") if __name__ == "__main__": main()