如何在M芯片的Macbook上训练神经网络

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/metal.html

import torch
# 检查是否可以使用 Metal GPU
print("Is Metal available?", torch.backends.mps.is_available())
# 检查 PyTorch 是否检测到了 GPU
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
print("Using device:", device)

# 创建模型
model = MyModel().to(device)
# 创建输入数据并转移到 GPU
input_data = torch.randn(64, 3, 224, 224).to(device)
# 进行前向传播
output = model(input_data)

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
# 设置设备
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
# 定义简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(3 * 224 * 224, 10)  # 假设输入为 3x224x224 的图像
    def forward(self, x):
        x = x.view(x.size(0), -1)  # 展平
        return self.fc(x)
# 数据加载
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])
train_dataset = datasets.FakeData(transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
# 初始化模型、损失函数和优化器
model = SimpleModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
# 训练循环
for epoch in range(10):
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)  # 将数据移动到 GPU
        optimizer.zero_grad()  # 清空梯度
        output = model(data)   # 前向传播
        loss = criterion(output, target)  # 计算损失
        loss.backward()        # 反向传播
        optimizer.step()       # 更新参数
    print(f'Epoch {epoch + 1}, Loss: {loss.item()}')
print("Training complete!")