Функция autocast

Функция autocast применяется в виде контекстного менеджера для автоматического подбора типов данных (FP16 или FP32) в процессе выполнения операций на CUDA. Это позволяет ускорить обучение и снизить потребление видеопамяти без потери качества модели. Функция доступна в модуле torch.cuda.amp.

Синтаксис

torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)

Параметр enabled управляет включением автоматического смешанного обучения. Параметр dtype задаёт тип данных для пониженной точности. Параметр cache_enabled включает кеширование автокаста для ускорения.

Пример

Давайте выполним простые операции в контексте автокаста:

import torch

# Создаём тензор на GPU
t = torch.tensor([1., 2., 3., 4., 5.], device='cuda')

# Включаем автокаст
with torch.cuda.amp.autocast():
	res = t * 2

print(res)

Результат выполнения кода:

tensor([2., 4., 6., 8., 10.], device='cuda:0')

Пример

Теперь выполним матричное умножение с автокастом:

import torch

torch.manual_seed(0)
t1 = torch.randn(3, 4, device='cuda')
t2 = torch.randn(4, 5, device='cuda')

with torch.cuda.amp.autocast():
	res = torch.matmul(t1, t2)

print(res.shape)
print(res.dtype)

Результат выполнения кода:

torch.Size([3, 5])
torch.float16

Пример

Продемонстрируем отключение автокаста для отдельных участков кода:

import torch

t = torch.tensor([1., 2., 3., 4., 5.], device='cuda')

with torch.cuda.amp.autocast():
	res1 = t * 2
	# Отключаем автокаст для этого блока
	with torch.cuda.amp.autocast(enabled=False):
		res2 = t * 3

print(res1.dtype)
print(res2.dtype)

Результат выполнения кода:

torch.float16
torch.float32

Смотрите также

функцию is_available,
которая проверяет доступность CUDA
функцию device_count,
которая возвращает количество доступных GPU
функцию current_device,
которая возвращает индекс текущего устройства
функцию synchronize,
которая синхронизирует все операции на устройстве