Метод to
Метод to класса Module применяется для преобразования параметров
и буферов модели к заданному типу данных или для перемещения модели на
указанное устройство. Метод работает in-place и возвращает сам объект
модели. Первым параметром можно передать тип данных или устройство.
Вторым параметром можно дополнительно указать тип данных, если первым
было передано устройство.
Синтаксис
model.to(dtype, [device])
model.to(device, [dtype])
model.to(device)
Пример
Давайте переместим модель на GPU, если он доступен:
import torch
import torch.nn as nn
model = nn.Linear(10, 5)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
print(next(model.parameters()).device)
Результат выполнения кода:
device(type='cuda', index=0)
Пример
Давайте преобразуем все параметры модели к типу float16:
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(10, 20),
nn.ReLU(),
nn.Linear(20, 5)
)
model.to(torch.float16)
print(next(model.parameters()).dtype)
Результат выполнения кода:
torch.float16
Пример
Давайте одновременно переместим модель на GPU и преобразуем к типу float16:
import torch
import torch.nn as nn
model = nn.Linear(10, 5)
model.to('cuda', torch.float16)
param = next(model.parameters())
print(param.device, param.dtype)
Результат выполнения кода:
device(type='cuda', index=0) torch.float16
Пример
Давайте переместим модель обратно на CPU:
import torch
import torch.nn as nn
model = nn.Linear(10, 5)
model.to('cpu')
print(next(model.parameters()).device)
Результат выполнения кода:
device(type='cpu')
Пример
Давайте преобразуем модель к типу float64 для более высокой точности вычислений:
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(5, 3),
nn.Linear(3, 1)
)
model.to(torch.float64)
print(next(model.parameters()).dtype)
Результат выполнения кода:
torch.float64
