Экскурсия по PyTorch

Привет!

Еще до конца мая у нас выйдет перевод книги Франсуа Шолле "Глубокое обучение на Python" (примеры с использованием библиотек Keras и Tensorflow). Не пропустите!

Но мы, естественно, смотрим в надвигающееся будущее и начинаем присматриваться к еще более инновационной библиотеке PyTorch. Сегодня вашему вниманию предлагается перевод статьи Питера Голдсборо, готового устроить вам долгую прогулку ознакомительную экскурсию по этой библиотеке. Под катом много и интересно.

Последние два года я всерьез занимался TensorFlow – писал статьи по этой библиотеке, выступал с лекциями о расширении ее бэкенда, либо использовал в моих собственных исследованиях, связанных с глубоким обучением. За этой работой я довольно хорошо усвоил, каковы сильные, а каковы слабые стороны TensorFlow – а также познакомился с конкретными архитектурными решениями, оставляющими поле для конкуренции. С таким багажом я недавно присоединился к команде PyTorch в отделе по исследованиям искусственного интеллекта в компании Facebook (FAIR) – пожалуй, в настоящее время это сильнейший конкурент TensorFlow. Сегодня PyTorch весьма популярна в исследовательском сообществе; почему – расскажу в следующих абзацах.

В этой статье я хочу дать экспресс-обзор библиотеки PyTorch, пояснить, ради чего она создавалась и познакомить вас с ее API.

Общая картина и философия

Для начала рассмотрим, что представляет собой PyTorch с фундаментальной точки зрения, какую модель программирования приходится применять, работая с ней, а также как она вписывается в экосистему современных инструментов глубокого обучения:

В сущности, PyTorch – это библиотека на Python, обеспечивающая тензорные вычисления с GPU-ускорением, подобно NumPy. Сверх этого PyTorch предлагает насыщенный API для решения прикладных задач, связанных с нейронными сетями.

PyTorch отличается от других фреймворков машинного обучения тем, что здесь не используются статические расчетные графы – определяемые заранее, сразу и окончательно – как в TensorFlow, Caffe2 или MXNet. Напротив, расчетные графы в PyTorch динамические и определяются на лету. Таким образом, при каждом вызове слоев в модели PyTorch динамически определяется новый расчетный граф. Этот граф создается имплицитно – то есть, библиотека сама записывает поток данных, идущих через программу, и связывает вызовы функций (узлы) вместе (посредством ребер) в расчетный граф.

Сравнение динамических и статических графов

Давайте подробнее разберемся, чем статические графы отличаются от динамических. В целом, в большинстве сред программирования при сложении двух переменных x и y, означающих числа, получается их суммарное значение (результат сложения). Например, на Python:

In [1]: x = 4
In [2]: y = 2
In [3]: x + y
Out[3]: 6

Но не в TensorFlow. В TensorFlow x и y будут не числами как таковыми, а описателями узлов графа, представляющих эти значения, но не содержащих их явно. Более того (что даже важнее), при сложении x и y получится не сумма этих чисел, а описатель расчетного графа, который даст искомое значение лишь после того, как будет выполнен:

In [1]: import tensorflow as tf
In [2]: x = tf.constant(4)
In [3]: y = tf.constant(2)
In [4]: x + y
Out[4]: <tf.Tensor 'add:0' shape=() dtype=int32>

В принципе, когда мы пишем код TensorFlow, это по факту не программирование, а метапрограммирование. Мы пишем программу (наш код), которая, в свою очередь, создает другую программу (расчетный граф TensorFlow). Естественно, первая модель программирования значительно проще второй. Гораздо удобнее говорить и рассуждать в контексте реальных феноменов, а не их представлений.

Важнейшее достоинство PyTorch заключается в том, что ее модель исполнения гораздо ближе к первой парадигме, чем ко второй. В основе своей PyTorch – это самый обычный Python с поддержкой тензорных вычислений (как и NumPy), но с GPU-ускорением тензорных операций и, что наиболее важно, со встроенным автоматическим дифференцированием (AD). Поскольку большинство современных алгоритмов машинного обучения серьезно зависят от типов данных из линейной алгебры (матриц и векторов) и используют градиентную информацию для уточнения оценок, двух этих столпов PyTorch достаточно, чтобы справиться со сколь угодно масштабными задачами машинного обучения.

Возвращаясь к разбору простого вышеприведенного случая, можно убедиться, что программирование с PyTorch по ощущению напоминает «естественный» Python:

In [1]: import torch
In [2]: x = torch.ones(1) * 4
In [3]: y = torch.ones(1) * 2
In [4]: x + y
Out[4]:
 6
[torch.FloatTensor of size 1]

PyTorch немного отличается от базовой логики программирования на Python в одном конкретном аспекте: библиотека записывает выполнение работающей программы. То есть, PyTorch тихонько “выслеживает”, какие операции вы совершаете над ее типами данных, и за кулисами – опять! – собирает расчетный граф. Такой расчетный граф нужен для автоматического дифференцирования, поскольку должен в обратном направлении проходить по цепочке операций, давшей результирующее значение, чтобы вычислить производные (для обратного автоматического дифференцирования). Серьезное отличие этого расчетного графа (вернее, способа сборки этого расчетного графа) от варианта из TensorFlow или MXNet заключается в том, что новый граф собирается «жадно», на лету, при интерпретации каждого фрагмента кода.

Напротив, в Tensorflow расчетный граф строится лишь однажды, за это отвечает метапрограмма (ваш код). Более того, тогда как PyTorch динамически обходит граф в обратном направлении всякий раз, когда вы запрашиваете производную значения, TensorFlow просто вставляет в граф дополнительные узлы, которые (неявно) вычисляют эти производные и интерпретируются точно как все остальные узлы. Здесь разница между динамическими и статическими графами проявляется особенно отчетливо.

Выбор, с какими расчетными графами работать – статическими или динамическими – серьезно упрощает процесс программирование в одном из этих окружений. Поток управления – это аспект, на котором особенно сказывается данный выбор. В окружении со статическими графами поток управления должен быть представлен на уровне графа в виде специализированных узлов. Например, в Tensorflow для обеспечения ветвления есть операция tf.cond(), принимающая в качестве ввода три подграфа: условный подграф и два подграфа для двух веток развития условия: if и else. Аналогично, циклы в графах Ternsorflow следует представлять как операции tf.while(), принимающие в качестве ввода condition и подграф body. В ситуации с динамическим графом все это упрощается. Поскольку графы при каждой интерпретации просматриваются из кода Python как есть, управление потоком можно нативно реализовать на языке, используя условия if и циклы while, как в любой другой программе. Таким образом, неуклюжий и путаный код Tensorflow:

import tensorflow as tf

x = tf.constant(2, shape=[2, 2])
w = tf.while_loop(
  lambda x: tf.reduce_sum(x) < 100,
  lambda x: tf.nn.relu(tf.square(x)),
  [x])

Превращается в естественный и понятный код PyTorch:

import torch.nn
from torch.autograd import Variable

x = Variable(torch.ones([2, 2]) * 2)
while x.sum() < 100:
    x = torch.nn.ReLU()(x**2)

Естественно, с точки зрения легкости программирования польза динамических графов этим далеко не ограничивается. Просто иметь возможность проверять промежуточные значения при помощи инструкций print (а не при помощи узлов tf.Print()) или в отладчике – уже большой плюс. Разумеется, динамизм может как оптимизировать программируемость, так и ухудшать производительность – то есть, оптимизировать такие графы сложнее. Поэтому, отличия и компромиссы между PyTorch и TensorFlow во многом такие же, как и между динамическим интерпретируемым языком, например, Python, и статическим компилируемым языком, например, C или C++. Первый проще и работать с ним быстрее, а из второго и третьего удобнее собирать сущности, хорошо поддающиеся оптимизации. Это и есть компромисс между гибкостью и производительностью.

Замечание об API PyTorch

Хочу сделать общее замечание по поводу API PyTorch, в особенности касающееся расчета нейронных сетей по сравнению с другими библиотеками, например, TensorFlow или MXNet — этот API обвешан множеством модулей (т.н. «batteries-included»). Как отметил один мой коллега, API Tensorflow так по-настоящему и не вышел за «сборочный» уровень, в том смысле, что этот API предоставляет лишь простейшие инструкции по сборке, необходимые для создания расчетных графов (сложение, умножение, поточечные функции, т.д.). Но он лишен «стандартной библиотеки» для наиболее распространенных программных фрагментов, которые программисту при работе приходится воспроизводить тысячи раз. Поэтому, чтобы выстраивать более высокоуровневые API поверх Tensorflow, приходится полагаться на помощь сообщества.
Действительно, сообщество создало такие высокоуровневые API. Правда, к сожалению, не один, а с десяток – в порядке соперничества. Таким образом, в неудачный день можно прочитать пять статей по своей специализации – и во всех пяти обнаружить разные «фронтенды» для TensorFlow. Как правило, между этими API совсем мало общего, так что вам по факту придется изучить 5 разных фреймворков, а не только TensorFlow. Вот некоторые наиболее популярные из этих API:

• Keras
• TFLearn
• PrettyTensor
• TF-Slim

PyTorch, в свою очередь, уже оснащена самыми ходовыми элементами, нужными для ежедневных исследований в области глубокого обучения. В принципе, в ней есть «нативный» Keras-подобный API в пакете torch.nn, обеспечивающий сцепление высокоуровневых модулей нейронных сетей.

Место PyTorch в общей экосистеме

Итак, объяснив, чем PyTorch отличается от статических графовых фреймворков вроде MXNet, TensorFlow или Theano, должен сказать, что PyTorch, фактически, не уникальна в своем подходе к вычислению нейронных сетей. До PyTorch уже существовали другие библиотеки, например, Chainer или DyNet, предоставлявшие подобный динамический API. Однако, сегодня PyTorch популярнее этих альтернатив.

Кроме того, PyTorch – не единственный фреймворк, используемый в Facebook. Основная рабочая нагрузка в продакшене у нас сейчас приходится на Caffe2 – это статический графовый фреймворк, выстроенный на основе Caffe. Чтобы подружить ту гибкость, что дает исследователю PyTorch, с достоинствами статических графов в сфере продакшен-оптимизации, в Facebook также разрабатывают ONNX, своеобразный формат обмена обмена информацией между PyTorch, Caffe2 и другими библиотеками, например, MXNet или CNTK.

Наконец, маленькое историческое отступление: до PyTorch, существовала Torch – совсем старая (образца начала 2000-х) библиотека для научных вычислений, написанная на языке Lua. Torch обертывает базу кода, написанную на C, благодаря чему она становится быстрой и эффективной. В принципе, PyTorch обертывает ровно ту же базу кода на C (правда, с дополнительным промежуточным уровнем абстрагирования), а пользователю выставляет API на Python. Далее поговорим об этом API на Python.

Работа с PyTorch

Далее мы обсудим базовые концепции и ключевые компоненты библиотеки PyTorch, изучим ее базовые типы данных, механизмы автоматического дифференцирования, специфический функционал, связанный с нейронными сетями, а также утилиты для загрузки и обработки данных.

Тензоры

Наиболее фундаментальный тип данных в PyTorch — это tensor. Тип данных tensor по значению и функциям очень похож на ndarray из NumPy. Более того, поскольку PyTorch нацелена на разумную интероперабельность с NumPy, API tensor также напоминает API ndarray (но не идентичен ему). Тензоры PyTorch можно создавать при помощи конструктора torch.Tensor, принимающего в качестве ввода размерности тензора и возвращающий тензор, который занимает неинициализированную область памяти:

import torch
x = torch.Tensor(4, 4)

На практике чаще всего приходится использовать одну из следующих функций PyTorch, возвращающих тензоры, инициализированные тем или иным образом, например:

• torch.rand: значения инициализируются из случайного равномерного распределения,
• torch.randn: значения инициализируются из случайного нормального распределения,
• torch.eye(n): единичная матрица вида n×nn×n,
• torch.from_numpy(ndarray): тензор PyTorch на основе ndarray из NumPy
• torch.linspace(start, end, steps): 1-D тензор со значениями steps, равномерно распределенными между start и end,
• torch.ones : тензор с одними единицами,
• torch.zeros_like(other): тензор такой же формы, что и other и с одними нулями,
• torch.arange(start, end, step): 1-D тензор со значениями, заполненными из диапазона.

Аналогично ndarray из NumPy, тензоры PyTorch предоставляют очень насыщенный API для комбинации с другими тензорами, а также для ситуативных изменений. Также, как и в NumPy, унарные и бинарные операции обычно можно выполнить при помощи функций из модуля torch, например, torch.add(x, y) или непосредственно при помощи методов в тензорных объектах, например, x.add(y). Для самых общих мест найдутся операторы перегрузки, например, x + y. Более того, для многих функций существуют ситуативные альтернативы, которые будут не создавать новый тензор, а изменять экземпляр получателя. Эти функции называются так же, как и стандартные варианты, однако, содержат в названии нижнее подчеркивание, например: x.add_(y).

Избранные операции:

torch.add(x, y): поэлементное сложение
torch.mm(x, y): умножение матриц (не matmul или dot),
torch.mul(x, y): поэлементное умножение
torch.exp(x): поэлементная экспонента
torch.pow(x, power): поэлементное возведение в степень
torch.sqrt(x): поэлементное возведение в квадрат
torch.sqrt_(x): ситуативное поэлементное возведение в квадрат
torch.sigmoid(x): поэлементная сигмоида
torch.cumprod(x): произведение всех значений
torch.sum(x): сумма всех значений
torch.std(x): стандартное отклонение всех значений
torch.mean(x): среднее всех значений

Тензоры во многом поддерживают семантику, знакомую по ndarray из NumPy, например, транслирование, сложное (прихотливое) индексирование (x[x > 5]) и поэлементные реляционные операторы (x > y). Тензоры PyTorch также можно преобразовывать непосредственно в ndarray NumPy при помощи функции torch.Tensor.numpy(). Наконец, поскольку основное превосходство тензоров PyTorch по сравнению с ndarray NumPy – это GPU-ускорение, к вашим услугам также есть функция torch.Tensor.cuda(), копирующая тензорную память на GPU-устройство с поддержкой CUDA, если таковое имеется.

Autograd

В центре большинства современных приемов машинного обучения лежит расчет градиентов. Это в особенности касается нейронных сетей, где для обновления весовых коэффициентов используется алгоритм обратного распространения. Именно поэтому в Pytorch есть сильная нативная поддержка градиентного вычисления функций и переменных, определенных внутри фреймворка. Такая техника, при которой градиенты автоматически рассчитываются для произвольных вычислений, называется автоматическим (иногда — алгоритмическим) дифференцированием.

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

Однако, в PyTorch нет статически вычисляемых графов, поэтому, здесь мы не можем позволить себе роскошь добавлять градиентные узлы уже после того, как определены остальные вычисления. Вместо этого PyTorch приходится записывать или прослеживать поток значений через программу по мере их поступления, то есть, динамически строить расчетный граф. Как только такой граф будет записан, у PyTorch будет информация, нужная для обратного обхода такого потока вычислений и расчета градиентов выходных значений на базе входных.

Tensor из PyTorch пока не обладает полноценными механизмами для участия в автоматическом дифференцировании. Чтобы тензор можно было записывать, его нужно обернуть в torch.autograd.Variable. Класс Variable предоставляет практически такой же API, как и Tensor, но дополняет его возможностью взаимодействия с torch.autograd.Function именно ради автоматического дифференцирования. Точнее, в Variable записывается история операций над Tensor.

Пользоваться torch.autograd.Variable очень просто. Нужно просто передать ему Tensor и сообщить torch, требует ли эта переменная записывать градиенты:

x = torch.autograd.Variable(torch.ones(4, 4), requires_grad=True)

Функция requires_grad может потребовать значения False, например, при вводе данных или работе с метками, поскольку такая информация обычно не дифференцируется. Однако, они все равно должны быть Variables, чтобы подходить для автоматического дифференцирования. Обратите внимание: requires_grad по умолчанию равна False, следовательно, для обучаемых параметров ее нужно устанавливать в True.

Для расчета градиентов и выполнения автоматического дифференцирования к Variable применяют функцию backward(). Так вычисляется градиент этого тензора относительно листьев расчетного графа (всех входных значений, повлиявших на данное). Затем эти градиенты собираются в член grad класса Variable:

In [1]: import torch
In [2]: from torch.autograd import Variable
In [3]: x = Variable(torch.ones(1, 5))
In [4]: w = Variable(torch.randn(5, 1), requires_grad=True)
In [5]: b = Variable(torch.randn(1), requires_grad=True)
In [6]: y = x.mm(w) + b # mm = matrix multiply
In [7]: y.backward() # perform automatic differentiation
In [8]: w.grad
Out[8]:
Variable containing:
 1
 1
 1
 1
 1
[torch.FloatTensor of size (5,1)]
In [9]: b.grad
Out[9]:
Variable containing:
 1
[torch.FloatTensor of size (1,)]
In [10]: x.grad
None

Поскольку все Variable кроме входных значений являются результатами операций, с каждой Variable ассоциирован grad_fn, представляющий собой функцию torch.autograd.Function для расчета обратного шага. Для входных значений он равен None:

In [11]: y.grad_fn
Out[11]: <AddBackward1 at 0x1077cef60>
In [12]: x.grad_fn
None
torch.nn

Модуль torch.nn предоставляет пользователям PyTorch функционал, специфичный для нейронных сетей. Один из важнейших его членов — torch.nn.Module, представляющий многоразовый блок операций и связанные с ним (обучаемые) параметры, чаще всего используемые в слоях нейронных сетей. Модули могут содержать иные модули и неявно получать функцию backward() для обратного распространения. Пример модуля — torch.nn.Linear(), представляющий линейный (плотный/полносвязный) слой (т.e. аффинное преобразование Wx+bWx+b):

In [1]: import torch
In [2]: from torch import nn
In [3]: from torch.autograd import Variable
In [4]: x = Variable(torch.ones(5, 5))
In [5]: x
Out[5]:
Variable containing:
 1  1  1  1  1
 1  1  1  1  1
 1  1  1  1  1
 1  1  1  1  1
 1  1  1  1  1
[torch.FloatTensor of size (5,5)]
In [6]: linear = nn.Linear(5, 1)
In [7]: linear(x)
Out[7]:
Variable containing:
 0.3324
 0.3324
 0.3324
 0.3324
 0.3324
[torch.FloatTensor of size (5,1)]

При обучении часто приходится вызывать в модуле функцию backward(), чтобы вычислять градиенты для его переменных. Поскольку при вызове backward() устанавливается член grad у Variables, также существует метод nn.Module.zero_grad(), сбрасывающий член grad всех Variable на ноль. Ваш обучающий цикл обычно вызывает zero_grad() в самом начале, либо непосредственно перед вызовом backward(), чтобы сбросить градиенты для следующего шага оптимизации.

При написании собственных моделей для нейронных сетей зачастую приходится писать собственные подклассы модуля для инкапсуляции распространенного функционала, который вы хотите интегрировать с PyTorch. Это делается очень просто – наследуем класс от torch.nn.Module и даем ему метод forward. Например, вот модуль, который я написал для одной из моих моделей (в ней ко входной информации добавляется гауссовский шум):

class AddNoise(torch.nn.Module):
    def __init__(self, mean=0.0, stddev=0.1):
        super(AddNoise, self).__init__()
        self.mean = mean
        self.stddev = stddev

    def forward(self, input):
        noise = input.clone().normal_(self.mean, self.stddev)
        return input + noise

Для соединения или сцепления модулей в полнофункциональные модели можно воспользоваться контейнером torch.nn.Sequential(), которому передают последовательность модулей – и он, в свою очередь, начинает действовать как самостоятельный модуль, при каждом вызове последовательно вычисляющий те модули, которые ему передали. Например:

In [1]: import torch
In [2]: from torch import nn
In [3]: from torch.autograd import Variable
In [4]: model = nn.Sequential(
   ...:     nn.Conv2d(1, 20, 5),
   ...:     nn.ReLU(),
   ...:     nn.Conv2d(20, 64, 5),
   ...:     nn.ReLU())
   ...:

In [5]: image = Variable(torch.rand(1, 1, 32, 32))
In [6]: model(image)
Out[6]:
Variable containing:
(0 ,0 ,.,.) =
  0.0026  0.0685  0.0000  ...   0.0000  0.1864  0.0413
  0.0000  0.0979  0.0119  ...   0.1637  0.0618  0.0000
  0.0000  0.0000  0.0000  ...   0.1289  0.1293  0.0000
           ...             ⋱             ...
  0.1006  0.1270  0.0723  ...   0.0000  0.1026  0.0000
  0.0000  0.0000  0.0574  ...   0.1491  0.0000  0.0191
  0.0150  0.0321  0.0000  ...   0.0204  0.0146  0.1724

Потери

torch.nn также предоставляет ряд функций потерь, естественно, важных для приложений в сфере машинного обучения. Примеры таких функций:

• torch.nn.MSELoss: средняя квадратичная функция потерь
• torch.nn.BCELoss: функция потерь бинарной кросс-энтропии,
• torch.nn.KLDivLoss: функция потерь информационного расхождения Кульбака-Лейблера

В контексте PyTorch функции потерь часто именуются критериями. В сущности, критерии – это очень простые модули, которые можно параметризовать непосредственно после создания, а с этого момента использовать как обычные функции:

In [1]: import torch
In [2]: import torch.nn
In [3]: from torch.autograd import Variable
In [4]: x = Variable(torch.randn(10, 3))
In [5]: y = Variable(torch.ones(10).type(torch.LongTensor))
In [6]: weights = Variable(torch.Tensor([0.2, 0.2, 0.6]))
In [7]: loss_function = torch.nn.CrossEntropyLoss(weight=weights)
In [8]: loss_value = loss_function(x, y)
Out [8]: Variable containing:
 1.2380
[torch.FloatTensor of size (1,)]

Оптимизаторы

После «первоэлементов» нейронных сетей (nn.Module) и функций потерь остается рассмотреть только оптимизатор, запускающий стохастический градиентный спуск (вариант). Для этого в PyTorch предоставляется пакет torch.optim, в котором определяется ряд распространенных алгоритмов оптимизации, в частности:

• torch.optim.SGD: стохастический градиентный спуск,
• torch.optim.Adam: адаптивная оценка моментов,
• torch.optim.RMSprop:алгоритм, разработанный Джеффри Хинтоном в рамках его курса для Coursera,
• torch.optim.LBFGS: алгоритм Бройдена-Флетчера-Гольдфарба-Шанно с ограниченным использованием памяти

Каждый из этих оптимизаторов создается со списком объектов-параметров, обычно извлекаемых методом parameters() из подкласса nn.Module, определяющим, какие значения будет обновлять оптимизатор. Кроме такого списка параметров каждый оптимизатор принимает некоторое количество дополнительных аргументов, помогающих сконфигурировать стратегию оптимизации. Например:

In [1]: import torch
In [2]: import torch.optim
In [3]: from torch.autograd import Variable
In [4]: x = Variable(torch.randn(5, 5))
In [5]: y = Variable(torch.randn(5, 5), requires_grad=True)
In [6]: z = x.mm(y).mean() # Perform an operation
In [7]: opt = torch.optim.Adam([y], lr=2e-4, betas=(0.5, 0.999))
In [8]: z.backward() # Calculate gradients
In [9]: y.data
Out[9]:
-0.4109 -0.0521  0.1481  1.9327  1.5276
-1.2396  0.0819 -1.3986 -0.0576  1.9694
 0.6252  0.7571 -2.2882 -0.1773  1.4825
 0.2634 -2.1945 -2.0998  0.7056  1.6744
 1.5266  1.7088  0.7706 -0.7874 -0.0161
[torch.FloatTensor of size 5x5]
In [10]: opt.step() # Обновляем y по правилам обновления градиентов Adam
In [11]: y.data
Out[11]:
-0.4107 -0.0519  0.1483  1.9329  1.5278
-1.2398  0.0817 -1.3988 -0.0578  1.9692
 0.6250  0.7569 -2.2884 -0.1775  1.4823
 0.2636 -2.1943 -2.0996  0.7058  1.6746
 1.5264  1.7086  0.7704 -0.7876 -0.0163
[torch.FloatTensor of size 5x5]

Загрузка данных

Для удобства в PyTorch предоставляется ряд утилит для загрузки датасетов, их предварительной обработки и взаимодействия с ними. Эти вспомогательные классы находятся в модуле torch.utils.data module. Здесь следует обратить внимание на две основные концепции:

1. Dataset, инкапсулирующий источник данных,
2. DataLoader, отвечающий за загрузку датасета, возможно, в параллельном режиме.

Для создания новых датасетов наследуется класс torch.utils.data.Dataset и переопределяется метод __len__, так, чтобы он возвращал количество образцов в датасете, а также метод __getitem__ для доступа к единичному значению по конкретному индексу. Например, так выглядит простой датасет, в котором инкапсулирован диапазон целых чисел:

import math

class RangeDataset(torch.utils.data.Dataset):
  def __init__(self, start, end, step=1):
    self.start = start
    self.end = end
    self.step = step

  def __len__(self, length):
    return math.ceil((self.end - self.start) / self.step)

  def __getitem__(self, index):
    value = self.start + index * self.step
    assert value < self.end
    return value

Внутри __init__ обычно конфигурируются какие-либо пути или изменяется набор возвращаемых в конечном итоге образцов. В __len__ указывается верхний предел индекса, с которым может быть вызван __getitem__, а в __getitem__ возвращается конкретный образец, например, изображение или аудиофрагмент.

Чтобы перебрать датасет, можно, в принципе, применить цикл for i in range и обращаться к образцам при помощи __getitem__. Однако, было бы гораздо удобнее, если бы датасет сам реализовывал протокол итератора, и мы могли бы сами перебирать образцы при помощи for sample in dataset. К счастью, такой функционал предоставляется в классе DataLoader. Объект DataLoader принимает датасет и ряд опций, конфигурирующих процедуру извлечения образца. Например, можно параллельно загружать образцы, задействовав множество процессов. Для этого конструктор DataLoader принимает аргумент num_workers. Обратите внимание: DataLoader всегда возвращает пакеты, размер которых задается в параметре batch_size. Простой пример:

dataset = RangeDataset(0, 10)
data_loader = torch.utils.data.DataLoader(
    dataset, batch_size=4, shuffle=True, num_workers=2, drop_last=True)

for i, batch in enumerate(data_loader):
  print(i, batch)

Здесь значение batch_size равно 4, поэтому возвращаемые тензоры будут содержать ровно по четыре значения. Если передать shuffle=True, то последовательность индексов для доступа к данным перемешивается, так что отдельные образцы возвращаются в случайном порядке. Мы также передали drop_last=True, поэтому, если для последнего пакета в датасете осталось меньше образцов, чем указано в batch_size, то этот пакет не возвращается. Наконец, мы задали для num_workers значение «два», то есть, выборкой данных параллельно займутся два процесса. После того, как DataLoader будет создан, перебор датасета и, соответственно, извлечение пакетов, станет простым и естественным.

Вот последнее интересное наблюдение, которым я хочу поделиться: DataLoader содержит довольно нетривиальную логику, определяющую, как комплектовать отдельные образцы, возвращенные в методе __getitem__ вашего датасета, в очередной пакет, возвращаемый DataLoader при переборе. Например, если __getitem__ возвращает словарь, то DataLoader агрегирует значения этого словаря в единое отображение, соответствующее одному пакету, использующему одинаковые ключи. Это значит, что, если метод __getitem__ датасета возвращает dict(example=example, label=label), то пакет, возвращенный DataLoader, вернет нечто наподобие dict(example=[example1, example2, ...], label=[label1, label2, ...]), то есть, распаковывая значения отдельных образцов, мы переупаковываем их под единым ключом для словаря пакета. Чтобы переопределить это поведение, можно передать аргумент функции для параметра collate_fn объекту DataLoader.

Обратите внимание: в пакете torchvision предоставляется ряд готовых датасетов, например, torchvision.datasets.CIFAR10. То же касается пакетов torchaudio и torchtext.

Заключение

Итак, теперь вы должны понимать и философию PyTorch, и ее базовый API, а значит, готовы перейти к покорению моделей PyTorch. Если ранее вы не сталкивались с PyTorch, но имеете опыт работы с другими фреймворками глубокого обучения, возьмите вашу любимую модель нейронной сети и перепишите ее при помощи PyTorch. Например, я переписал для PyTorch архитектуру LSGAN, реализованную для TensorFlow, и при этом изрядно с ней напрактиковался. Вас также могут заинтересовать статьи, опубликованные здесь и здесь.

Автор: ph_piter

Источник