Principio de aceleración
Apple tiene su propia API de implementación de GPU: Metal, y la aceleración de Pytorch se basa en Metal. Específicamente, el uso de Metal Performance Shaders (MPS) de Apple como backend de PyTorch puede acelerar el entrenamiento de GPU. El backend de MPS amplía el marco PyTorch y proporciona scripts y funciones para configurar y ejecutar operaciones en la Mac. MPS optimiza el rendimiento informático con núcleos ajustados para las características únicas de cada familia de GPU Metal. El nuevo dispositivo mapea gráficos y primitivos computacionales de aprendizaje automático sobre el marco de trabajo de gráficos MPS y los núcleos ajustados proporcionados por MPS.
Por lo tanto, el nombre del dispositivo recién agregado es mps y el método de uso es similar a cuda, por ejemplo:
import torch
foo = torch.rand(1, 3, 224, 224).to('mps')
device = torch.device('mps')
foo = foo.to(device)
Además, descubrí que Pytorch ya es compatible con los siguientes dispositivos, lo cual es realmente inesperado:
cpu, cuda, ipu, xpu, mkldnn, opengl, opencl,
ideep, hip, ve, ort, mps, xla, lazy, vulkan, meta, hpu
Configuración del entorno
Para utilizar esta función experimental, debe cumplir las siguientes tres condiciones:
-
Tener una computadora portátil Mac con chips Apple Silicon (M1, M1 Pro, M1 Pro Max, M1 Ultra)
-
ArmPython de 64 bits instalado
-
Instalada la última versión nocturna de Pytorch
Suponiendo que la máquina esté lista. Podemos descargar la versión arm64 de miniconda desde aquí (el nombre del archivo es Miniconda3 macOS Apple M1 64-bit bash, y el entorno Python instalado en base a él es arm64. Los comandos para descargar e instalar Minicoda son los siguientes:
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh
chmod +x Miniconda3-latest-MacOSX-arm64.sh
./Miniconda3-latest-MacOSX-arm64.sh
Simplemente siga las instrucciones. Una vez completada la instalación, cree un entorno virtual y compruebe la arquitectura de Python comprobando si platform.uname()[4] es arm64:
conda config --env --set always_yes true
conda create -n try-mps python=3.8
conda activate try-mps
python -c "import platform; print(platform.uname()[4])"
Si el resultado del último comando es arm64, significa que la versión de Python está bien y puede continuar.
El tercer paso es instalar la versión nocturna de Pytorch y realizar las siguientes operaciones en el entorno virtual abierto:
python -m pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
Una vez completada la ejecución, verifique si el backend MPS está disponible a través del siguiente comando:
python -c "import torch;print(torch.backends.mps.is_built())"
Si el resultado es True, significa que el backend de MPS está disponible y puede continuar.
ejecutar un MNIST
Basado en el ejemplo MNIST en el ejemplo oficial de Pytorch, se ha modificado para probar los modos cpu y mps. El código es el siguiente:
from __future__ import print_function
import argparse
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout1 = nn.Dropout(0.25)
self.dropout2 = nn.Dropout(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
output = F.log_softmax(x, dim=1)
return output
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
if args.dry_run:
break
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--batch-size', type=int, default=64, metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--epochs', type=int, default=1, metavar='N',
help='number of epochs to train (default: 14)')
parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
help='learning rate (default: 1.0)')
parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--use_gpu', action='store_true', default=False,
help='enable MPS')
parser.add_argument('--dry-run', action='store_true', default=False,
help='quickly check a single pass')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=False,
help='For Saving the current Model')
args = parser.parse_args()
use_gpu = args.use_gpu
torch.manual_seed(args.seed)
device = torch.device("mps" if args.use_gpu else "cpu")
train_kwargs = {
'batch_size': args.batch_size}
if use_gpu:
cuda_kwargs = {
'num_workers': 1,
'pin_memory': True,
'shuffle': True}
train_kwargs.update(cuda_kwargs)
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
dataset1 = datasets.MNIST('../data', train=True, download=True,
transform=transform)
dataset2 = datasets.MNIST('../data', train=False,
transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
scheduler.step()
if __name__ == '__main__':
t0 = time.time()
main()
t1 = time.time()
print('time_cost:', t1 - t0)
CPU de prueba:
python main.py
MPS de prueba:
python main --use_gpu
La prueba en mi máquina M1 encontró que entrenar una Epoch de MNIST tomó 149.6 s de tiempo de CPU, mientras que usar MPS tomó 18.4 s. El efecto de mejora es significativo, o puede ser que la CPU esté funcionando demasiado rápido. En una palabra, si puede usar mps para entrenar el modelo, debe usar mps. La CPU es demasiado lenta.