Новый PyTorch получил стабильный модуль FFT

Богдан Мирченко

16.03.2021 17:42

Состоялся релиз фреймворка машинного обучения PyTorch 1.8. В новой версии улучшены поддержка быстрого преобразования Фурье (FFT) и обучение распределенной модели, появились новые API-интерфейсы, обновились библиотеки и многое другое.

Наиболее важным обновлением можно назвать стабильный модуль torch.fft. Он поддерживает высокопроизводительные быстрые преобразования Фурье в стиле модуля np.fft в NumPy, но с поддержкой аппаратного ускорения и автоградацией. В предыдущих версиях torch.fft был отдельной функцией.

Другие улучшения API включают новые функции линейной алгебры, поддержку автоградации для сложных тензоров и повышение производительности для вычисления гессиан функций и якобианов. Разработчики заявляют, что значительные обновления и улучшения внесены в распределенное обучение. Эти повышает надежность NCCL, добавляет поддержку конвейерного параллелизма и профилирования RPC.

В режиме бета Python 1.8 получил утилиты для тестирования производительности Benchmark. Они позволяют проводить точные измерения и дают инструменты, помогающие с анализом тестов и их последующей обработкой. Создатели фреймворка надеются, что это поможет разработчикам PyTorch быстро понять, как их действия влияют на производительность.

Например:

from torch.utils.benchmark import Timer

results = []
for num_threads in [1, 2, 4]:
    timer = Timer(
        stmt="torch.add(x, y, out=out)",
        setup="""
            n = 1024
            x = torch.ones((n, n))
            y = torch.ones((n, 1))
            out = torch.empty((n, n))
        """,
        num_threads=num_threads,
    )
    results.append(timer.blocked_autorange(min_run_time=5))
    print(
        f"{num_threads} thread{'s' if num_threads > 1 else ' ':<4}"
        f"{results[-1].median * 1e6:>4.0f} us   " +
        (f"({results[0].median / results[-1].median:.1f}x)" if num_threads > 1 else '')
    )

1 thread     376 us   
2 threads    189 us   (2.0x)
4 threads     99 us   (3.8x)

Разработчики обновили ряд функций для улучшения совместимости с NumPy. К ним, например, относится модуль torch.linalg, созданный по образцу модуля np.linalg. Он обеспечивает поддержку общих операций линейной алгебры в стиле NumPy, в том числе — разложение Холецкого и детерминанты.

С новой версией команда PyTorch также выпустила основные обновления библиотек, включая TorchSPRNG, TorchVision, TorchText и TorchAudio.