Метод half
Метод half преобразует все параметры модели (веса и смещения)
и буферы в тип данных float16. Это позволяет значительно
ускорить вычисления на современных GPU и сократить использование
видеопамяти. Метод действует на месте (in-place) и возвращает
сам модуль, что удобно для цепочечных вызовов.
Синтаксис
model.half()
Пример
Давайте создадим простую модель и преобразуем её веса в тип float16⁅/⁆:
import torch
import torch.nn as nn
model = nn.Linear(10, 5)
print(model.weight.dtype)
model = model.half()
print(model.weight.dtype)
Результат выполнения кода:
torch.float32
torch.float16
Пример
Давайте выполним прямой проход через модель после преобразования в float16⁅/⁆:
import torch
import torch.nn as nn
torch.manual_seed(0)
model = nn.Linear(10, 5)
model = model.half()
x = torch.randn(2, 10).half()
out = model(x)
print(out.dtype)
print(out)
Результат выполнения кода:
torch.float16
tensor([
[ 0.4624, -0.2170, -0.3408, 0.3076, -0.4336],
[ 0.1788, -0.2156, -0.4312, 0.0128, 0.1599]
], dtype=torch.float16)
Обратите внимание, что входной тензор также должен быть преобразован
в тип float16⁅/⁆ методом half⁅/⁆ перед передачей в модель.
Пример
Давайте сравним потребление памяти до и после преобразования модели:
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(1000, 1000),
nn.ReLU(),
nn.Linear(1000, 1000)
)
param_size = sum(p.numel() for p in model.parameters())
memory_float32 = param_size * 4 / (1024 ** 2)
print(f"float32: {memory_float32:.2f} MB")
model = model.half()
memory_float16 = param_size * 2 / (1024 ** 2)
print(f"float16: {memory_float16:.2f} MB")
Результат выполнения кода:
"float32: 7.63 MB"
"float16: 3.81 MB"
Пример
Давайте используем метод half⁅/⁆ в сочетании с другими методами преобразования:
import torch
import torch.nn as nn
model = nn.Linear(10, 5)
model = model.half().cuda()
print(model.weight.dtype)
print(next(model.parameters()).device)
Результат выполнения кода:
torch.float16
"cuda:0"
