Метод cuda
Метод cuda класса Module переносит все параметры
и постоянные буферы модели на указанное устройство CUDA.
Метод принимает необязательный параметр device -
индекс устройства GPU, на которое будет перенесена модель.
Если параметр не указан, используется устройство с индексом 0.
Метод возвращает сам объект модели, что позволяет
использовать цепочки вызовов.
Синтаксис
model.cuda([device])
Пример
Давайте создадим простую модель и перенесём её на GPU:
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(10, 20),
nn.ReLU(),
nn.Linear(20, 5),
)
print(model)
model = model.cuda()
Результат выполнения кода:
Sequential(
(0): Linear(in_features=10, out_features=20, bias=True)
(1): ReLU()
(2): Linear(in_features=20, out_features=5, bias=True)
)
Пример
Теперь передадим данные в модель и выполним прямой проход:
import torch
import torch.nn as nn
torch.manual_seed(0)
model = nn.Linear(5, 3)
model = model.cuda()
x = torch.randn(2, 5).cuda()
res = model(x)
print(res.device)
Результат выполнения кода:
cuda:0
Пример
Укажем конкретное устройство для переноса модели:
import torch
import torch.nn as nn
model = nn.Linear(5, 3)
model = model.cuda(1)
print(next(model.parameters()).device)
Результат выполнения кода:
cuda:1
Пример
Метод cuda можно использовать в цепочке
с другими методами, например, с to:
import torch
import torch.nn as nn
torch.manual_seed(0)
model = nn.Linear(5, 3)
model = model.cuda().to(torch.float16)
x = torch.randn(2, 5).cuda()
res = model(x)
print(res.dtype)
Результат выполнения кода:
torch.float16
