Сравнение матричных арифметических вычислений в CPU и GPU с Python и PyTorch

Как быстро прирост вычислений графического процессора сравнивается с процессором? В этой статье я собираюсь проверить это, используя функции линейного преобразования Python и PyTorch.

Вот некоторые характеристики моей тестовой машины:

  • Процессор: Intel i7 6700k (4 ядра/8 потоков)
  • Графический процессор: RTX 3070 TI (6144 ядра CUDA и 192 ядра Tensor)
  • Оперативная память: 32G
  • ОС: Виндовс 10

Объяснение жаргона графических процессоров NVIDIA

CUDA – это сокращение от Compute Unified Device Architecture. Вы можете использовать CUDA для прямого доступа к набору инструкций графического процессора NVIDIA.

В отличие от DirectX и OpenGL, специально разработанных для создания игровых движков, CUDA не требует от пользователей знания сложного языка графического программирования¹.

Тензорные ядра — это процессоры, которые ускоряют процесс умножения матриц².

Например, умножение двух матриц 4×4 с использованием CPU или CUDA включает 64 умножения и 48 сложений, одну операцию за такт, тогда как тензорные ядра могут выполнять несколько операций за такт.

Больше знакомства с тензорными ядрами в этом видео с канала Nvidia Developer на YouTube.

Какова связь между ядрами CUDA и ядрами Tensor? Тензорные ядра встроены в ядра CUDA, эти волшебные ядра будут запускаться при выполнении определенных условий.

Методология испытаний

Вычисления на GPU быстрее, чем на CPU, только в некоторых типичных сценариях. В других случаях вычисления на GPU могут быть медленнее, чем на CPU!

CUDA широко используется в машинном обучении и глубоком обучении из-за его особого качества при параллельном умножении и сложении матриц.

В математическом уравнении:

Линейная функция PyTorch torch.nn.Linear выполняет точно такую ​​же операцию. Например, вы можете преобразовать матрицу 2x2 в матрицу 2x3 с помощью следующего кода:

import torch
in_row,in_f,out_f = 2,2,3
tensor            = torch.randn(in_row,in_f)
l_trans           = torch.nn.Linear(in_f,out_f)
print(l_trans(tensor))

Базовый уровень ЦП

Прежде чем измерять производительность графического процессора, мне нужно установить базовую производительность центрального процессора.

Чтобы немного нагрузить микросхемы и продлить время работы, я увеличиваю числа in_row, in_f и out_f, также устанавливаю цикл работы 10 000  раз.

import torch
import torch.nn
import time
in_row, in_f, out_f = 256, 1024, 2048
loop_times = 10000

Теперь давайте посмотрим, сколько секунд потребуется процессору, чтобы завершить 10 000 преобразований:

s       = time.time()
tensor  = torch.randn(in_row, in_f).to('cpu')
l_trans = torch.nn.Linear(in_f, out_f).to('cpu')
for _ in range(loop_times):
    l_trans(tensor)
print('cpu take time:',time.time()-s)

Результат:

cpu take time: 55.70971965789795

Мой i7 6700k занимает около 55 секунд, если честно результат не плохой.

Вычислите это в GPU

Чтобы попросить CUDA графического процессора выполнить те же вычисления, я просто заменяю .to(‘cpu’) на .cude(). Кроме того, учитывая, что операции в CUDA являются асинхронными, мне также необходимо добавить оператор синхронизации, чтобы обеспечить печать использованного времени после выполнения всех задач CUDA.

s       = time.time()
tensor  = torch.randn(in_row, in_f).cuda()
l_trans = torch.nn.Linear(in_f, out_f).cuda()
for _ in range(loop_times):
    l_trans(tensor)
torch.cuda.synchronize()
print('CUDA take time:',time.time()-s)

Изменения кода выделены, вот текущий результат:

CUDA take time: 1.327127456665039

Почти в 42x раза быстрее, чем при работе на ЦП. Модель, которой требуется несколько дней обучения на CPU, теперь может занять всего несколько часов на GPU. Это очень быстро.

Включить тензорные ядра

CUDA уже работает быстро, как насчет включения 197 тензорных ядер RTX 3070Ti? Согласно этому видео, в PyTorch, чтобы включить тензорные ядра, все, что мне нужно сделать, это уменьшить точность с плавающей запятой с FP32 до FP16.

s       = time.time()
tensor  = torch.randn(in_row, in_f).cuda().half()
layer   = torch.nn.Linear(in_f, out_f).cuda().half()
for _ in range(loop_times):
    layer(tensor)
torch.cuda.synchronize()
print('CUDA with tensor cores take time:',time.time()-s)

Результат:

CUDA with tensor cores take time:0.5381264686584473

Еще одно улучшение производительности в 2,6 раза.

Заключение

В этой статье я сравнил операцию линейного преобразования, вызвав функцию линейного преобразования PyTorch в CPU, GPU CUDA и GPU CUDA + Tensor Cores. Вот обобщенный результат:

CUDA и тензорные ядра NVIDIA действительно значительно улучшают производительность умножения матриц.

Справочные ссылки

  1. УСКОРЕНИЕ ЧИСЛЕННЫХ РАСЧЕТОВ С ИСПОЛЬЗОВАНИЕМ ГРАФИЧЕСКОГО ПРОЦЕССОРА (GPU)
  2. Объяснение ядер Nvidia CUDA: чем они отличаются?
  3. Серия видео: методы обучения смешанной точности с использованием тензорных ядер для глубокого обучения