Esta entrada es parte del curso de Deep learning con PyTorch.
En un post pasado vimos como las GPUs (graphics processing units) han impulsado el desarrollo del Deep Learning en los últimos años.
Las GPUs, gracias a su estructura en paralelo, son muy eficientes para algoritmos con procesamiento de datos en paralelo como los usados en deep learning.
Para impulsar este nuevo paradigma, NVIDIA desarrolló la arquitectura de cálculo paralelo CUDA, que se incluye en sus GPUs (GeForce, ION Quadro y Tesla). De esta forma, miles de programadores pueden implementar este paralelismo en el procesamiento de tareas y datos usando diferentes lenguajes de programación de alto nivel.
En este post vamos a ver como de manera sencilla se puede usar CUDA para realizar computaciones usando la GPU en Google Colab y con PyTorch.
La librería PyTorch ya viene preinstalada en Colab junto con el soporte para CUDA. Con el método “torch.cuda.is_available()” se determina si una GPU está disponible.
Con el método “torch.cuda.get_device_name()” se consigue el nombre de la GPU utilizada.
Con “torch.cuda.memory_allocated()” se obtiene la cantidad de memoria máxima de la GPU usada por los tensores.
Para que el modelo se ejecute en la GPU basta con copiar los parámetros del modelo en la GPU con net.cuda().
A continuación podemos ver un ejemplo de una red neuronal con una capa intermedia en Pytorch para clasificar los dígitos MNIST ejecutada en la GPU.
La salida de este ejemplo, entre otras cosas, nos indicaría si tenemos una GPU disponible, qué tipo de GPU es y la memoria máxima utilizada:
“Using device: cuda”
“Tesla K80”
“Allocated: 6703104”
Como se puede ver la GPU disponible en Colab es la Tesla K80 de NVIDIA.
A continuación mostramos el código ejemplo (puede haber problemas con la sangría por el copiado del código).
import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 1e-3
train_dataset = dsets.MNIST(root='./data',
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = dsets.MNIST(root='./data',
train=False,
transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
class Net(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(Net, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out
net = Net(input_size, hidden_size, num_classes)
use_cuda = True
net.cuda()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
images = Variable(images.view(-1, 28*28))
labels = Variable(labels)
images = images.cuda()
labels = labels.cuda()
optimizer.zero_grad()
outputs = net(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if (i+1) % 200 == 0:
print('Epoch [%d/%d], Step [%d/%d], Loss: %.4f'
%(epoch+1, num_epochs, i+1, len(train_dataset)//batch_size, loss.data))
if device.type == 'cuda':
print(torch.cuda.get_device_name(0))
print('Allocated:', torch.cuda.memory_allocated(0))
Deja una respuesta