Функция cuda.synchronize

Функция torch.cuda.synchronize используется для синхронизации хоста и устройства CUDA. Она принудительно блокирует выполнение программы на CPU до тех пор, пока все ранее запущенные на GPU операции не будут полностью завершены. Это необходимо для корректных замеров времени, отладки и гарантии того, что данные из GPU памяти доступны для чтения на CPU. Функция принимает необязательный параметр device, указывающий устройство для синхронизации.

Если параметр device не указан, синхронизация выполняется для текущего устройства. В случае отсутствия доступных устройств CUDA вызов функции не приводит к ошибке, но и не выполняет никаких действий.

Синтаксис

torch.cuda.synchronize([device])

Пример

Давайте выполним синхронизацию после запуска нескольких операций на GPU:

import torch

t = torch.tensor([1, 2, 3, 4, 5], device='cuda')
t = t * 2
t = t + 10
torch.cuda.synchronize()
print(t)

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

tensor([12, 14, 16, 18, 20], device='cuda:0')

Пример

Пример синхронизации с явным указанием устройства для замера времени выполнения операций:

import torch
import time

torch.manual_seed(0)
t = torch.randn(1000, 1000, device='cuda')
start = time.time()
t = t @ t
torch.cuda.synchronize(device='cuda')
end = time.time()
print(end - start)

Результат выполнения кода (зависит от производительности GPU):

"0.003483"

Пример

Пример использования синхронизации для проверки корректности данных после операций на GPU:

import torch

t = torch.tensor([1, 2, 3, 4, 5], device='cuda')
t = t ** 2
torch.cuda.synchronize()
# После синхронизации можно безопасно копировать данные на CPU
res = t.cpu()
print(res)

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

tensor([ 1, 4, 9, 16, 25])

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

функцию is_available,
которая проверяет доступность CUDA
функцию current_device,
которая возвращает индекс текущего устройства
функцию empty_cache,
которая освобождает кэш памяти GPU
функцию memory_allocated,
которая возвращает объем выделенной памяти на GPU