sonnet
v2.0.2
Документация | Примеры
Sonnet — это библиотека, созданная на основе TensorFlow 2, предназначенная для предоставления простых компонуемых абстракций для исследований в области машинного обучения.
Sonnet был спроектирован и создан исследователями из DeepMind. Его можно использовать для построения нейронных сетей для самых разных целей (обучение без учителя, обучение с подкреплением и т. д.). Мы считаем, что это удачная абстракция для нашей организации, и вы тоже можете это сделать!
Точнее, Sonnet предоставляет простую, но мощную модель программирования, основанную на одной концепции: snt.Module . Модули могут содержать ссылки на параметры, другие модули и методы, которые применяют некоторые функции к пользовательскому вводу. Sonnet поставляется со многими предопределенными модулями (например, snt.Linear , snt.Conv2D , snt.BatchNorm ) и некоторыми предопределенными сетями модулей (например, snt.nets.MLP ), но пользователям также предлагается создавать свои собственные модули.
В отличие от многих фреймворков, Sonnet крайне непредвзято относится к тому, как вы будете использовать свои модули. Модули спроектированы так, чтобы быть автономными и полностью отделенными друг от друга. Sonnet не поставляется с обучающей структурой, и пользователям рекомендуется создавать свои собственные или использовать созданные другими.
Sonnet также спроектирован так, чтобы его было легко понять, наш код (надеюсь!) ясен и сфокусирован. Там, где мы выбрали значения по умолчанию (например, значения по умолчанию для начальных значений параметров), мы пытаемся указать, почему.
Самый простой способ попробовать Sonnet — использовать Google Colab, который предлагает бесплатный блокнот Python, подключенный к графическому процессору или TPU.
snt.distribute Для начала установите TensorFlow 2.0 и Sonnet 2:
$ pip install tensorflow tensorflow-probability
$ pip install dm-sonnetВы можете запустить следующее, чтобы проверить правильность установки:
import tensorflow as tf
import sonnet as snt
print ( "TensorFlow version {}" . format ( tf . __version__ ))
print ( "Sonnet version {}" . format ( snt . __version__ )) Sonnet поставляется с рядом встроенных модулей, которые вы можете легко использовать. Например, чтобы определить MLP, мы можем использовать модуль snt.Sequential для вызова последовательности модулей, передавая выходные данные данного модуля в качестве входных данных для следующего модуля. Мы можем использовать snt.Linear и tf.nn.relu , чтобы фактически определить наши вычисления:
mlp = snt . Sequential ([
snt . Linear ( 1024 ),
tf . nn . relu ,
snt . Linear ( 10 ),
]) Чтобы использовать наш модуль, нам нужно его «вызвать». Модуль Sequential (и большинство модулей) определяет метод __call__ , что означает, что вы можете вызывать их по имени:
logits = mlp ( tf . random . normal ([ batch_size , input_size ]))Также очень часто запрашиваются все параметры вашего модуля. Большинство модулей в Sonnet создают свои параметры при первом вызове с каким-либо вводом (поскольку в большинстве случаев форма параметров является функцией ввода). Модули Sonnet предоставляют два свойства для доступа к параметрам.
Свойство variables возвращает все tf.Variable , на которые ссылается данный модуль:
all_variables = mlp . variables Стоит отметить, что tf.Variable используются не только для параметров вашей модели. Например, они используются для хранения состояния в метриках, используемых в snt.BatchNorm . В большинстве случаев пользователи извлекают переменные модуля, чтобы передать их оптимизатору для обновления. В этом случае необучаемые переменные обычно не должны находиться в этом списке, поскольку они обновляются с помощью другого механизма. TensorFlow имеет встроенный механизм для маркировки переменных как «обучаемых» (параметры вашей модели) и необучаемых (другие переменные). Sonnet предоставляет механизм для сбора всех обучаемых переменных из вашего модуля, который, вероятно, вы и хотите передать оптимизатору:
model_parameters = mlp . trainable_variables Sonnet настоятельно рекомендует пользователям создавать подклассы snt.Module для определения своих собственных модулей. Начнем с создания простого Linear слоя под названием MyLinear :
class MyLinear ( snt . Module ):
def __init__ ( self , output_size , name = None ):
super ( MyLinear , self ). __init__ ( name = name )
self . output_size = output_size
@ snt . once
def _initialize ( self , x ):
initial_w = tf . random . normal ([ x . shape [ 1 ], self . output_size ])
self . w = tf . Variable ( initial_w , name = "w" )
self . b = tf . Variable ( tf . zeros ([ self . output_size ]), name = "b" )
def __call__ ( self , x ):
self . _initialize ( x )
return tf . matmul ( x , self . w ) + self . bИспользование этого модуля тривиально:
mod = MyLinear ( 32 )
mod ( tf . ones ([ batch_size , input_size ])) Создав подкласс snt.Module вы бесплатно получите множество приятных свойств. Например, реализация __repr__ по умолчанию, которая показывает аргументы конструктора (очень полезно для отладки и самоанализа):
>> > print ( repr ( mod ))
MyLinear ( output_size = 10 ) Вы также получаете variables и свойства trainable_variables :
>> > mod . variables
( < tf . Variable 'my_linear/b:0' shape = ( 10 ,) ...) > ,
< tf . Variable 'my_linear/w:0' shape = ( 1 , 10 ) ...) > ) Вы можете заметить префикс my_linear в переменных выше. Это связано с тем, что модули Sonnet также входят в область имен модулей при каждом вызове методов. Вводя область имени модуля, мы предоставляем гораздо более полезный граф для использования такими инструментами, как TensorBoard (например, все операции, которые происходят внутри my_linear, будут находиться в группе под названием my_linear).
Кроме того, ваш модуль теперь будет поддерживать контрольные точки TensorFlow и сохраненную модель, которые являются расширенными функциями, которые будут рассмотрены позже.
Sonnet поддерживает несколько форматов сериализации. Самый простой формат, который мы поддерживаем, — это pickle Python, и все встроенные модули тестируются, чтобы убедиться, что их можно сохранить/загрузить с помощью Pickle в одном и том же процессе Python. В целом мы не рекомендуем использовать Pickle, он плохо поддерживается многими частями TensorFlow и, по нашему опыту, может быть весьма хрупким.
Ссылка: https://www.tensorflow.org/alpha/guide/checkpoints.
Контрольные точки TensorFlow можно использовать для периодического сохранения значений параметров во время обучения. Это может быть полезно для сохранения прогресса обучения в случае сбоя или остановки вашей программы. Sonnet предназначен для корректной работы с контрольными точками TensorFlow:
checkpoint_root = "/tmp/checkpoints"
checkpoint_name = "example"
save_prefix = os . path . join ( checkpoint_root , checkpoint_name )
my_module = create_my_sonnet_module () # Can be anything extending snt.Module.
# A `Checkpoint` object manages checkpointing of the TensorFlow state associated
# with the objects passed to it's constructor. Note that Checkpoint supports
# restore on create, meaning that the variables of `my_module` do **not** need
# to be created before you restore from a checkpoint (their value will be
# restored when they are created).
checkpoint = tf . train . Checkpoint ( module = my_module )
# Most training scripts will want to restore from a checkpoint if one exists. This
# would be the case if you interrupted your training (e.g. to use your GPU for
# something else, or in a cloud environment if your instance is preempted).
latest = tf . train . latest_checkpoint ( checkpoint_root )
if latest is not None :
checkpoint . restore ( latest )
for step_num in range ( num_steps ):
train ( my_module )
# During training we will occasionally save the values of weights. Note that
# this is a blocking call and can be slow (typically we are writing to the
# slowest storage on the machine). If you have a more reliable setup it might be
# appropriate to save less frequently.
if step_num and not step_num % 1000 :
checkpoint . save ( save_prefix )
# Make sure to save your final values!!
checkpoint . save ( save_prefix )Ссылка: https://www.tensorflow.org/alpha/guide/saved_model.
Сохраненные модели TensorFlow можно использовать для сохранения копии вашей сети, отделенной для нее от исходного кода Python. Это становится возможным благодаря сохранению графа TensorFlow, описывающего вычисления, и контрольной точки, содержащей значения весов.
Первое, что нужно сделать для создания сохраненной модели, — это создать snt.Module , который вы хотите сохранить:
my_module = snt . nets . MLP ([ 1024 , 1024 , 10 ])
my_module ( tf . ones ([ 1 , input_size ]))Далее нам нужно создать еще один модуль, описывающий конкретные части нашей модели, которые мы хотим экспортировать. Мы советуем делать это (вместо изменения исходной модели на месте), чтобы иметь детальный контроль над тем, что на самом деле экспортируется. Обычно это важно, чтобы избежать создания очень больших сохраненных моделей и чтобы вы делились только теми частями вашей модели, которые хотите (например, вы хотите поделиться только генератором для GAN, но оставить дискриминатор закрытым).
@ tf . function ( input_signature = [ tf . TensorSpec ([ None , input_size ])])
def inference ( x ):
return my_module ( x )
to_save = snt . Module ()
to_save . inference = inference
to_save . all_variables = list ( my_module . variables )
tf . saved_model . save ( to_save , "/tmp/example_saved_model" ) Теперь у нас есть сохраненная модель в папке /tmp/example_saved_model :
$ ls -lh /tmp/example_saved_model
total 24K
drwxrwsr-t 2 tomhennigan 154432098 4.0K Apr 28 00:14 assets
-rw-rw-r-- 1 tomhennigan 154432098 14K Apr 28 00:15 saved_model.pb
drwxrwsr-t 2 tomhennigan 154432098 4.0K Apr 28 00:15 variablesЗагрузка этой модели проста и может быть выполнена на другом компьютере без какого-либо кода Python, создавшего сохраненную модель:
loaded = tf . saved_model . load ( "/tmp/example_saved_model" )
# Use the inference method. Note this doesn't run the Python code from `to_save`
# but instead uses the TensorFlow Graph that is part of the saved model.
loaded . inference ( tf . ones ([ 1 , input_size ]))
# The all_variables property can be used to retrieve the restored variables.
assert len ( loaded . all_variables ) > 0 Обратите внимание, что загруженный объект не является модулем Sonnet, это объект-контейнер, который имеет определенные методы (например, inference ) и свойства (например, all_variables ), которые мы добавили в предыдущем блоке.
Пример: https://github.com/deepmind/sonnet/blob/v2/examples/distributed_cifar10.ipynb
Sonnet поддерживает распределенное обучение с использованием пользовательских стратегий распространения TensorFlow.
Ключевое различие между Sonnet и распределенным обучением с использованием tf.keras заключается в том, что модули и оптимизаторы Sonnet не ведут себя по-разному при запуске в рамках стратегий распределения (например, мы не усредняем ваши градиенты и не синхронизируем статистику ваших пакетных норм). Мы считаем, что пользователи должны полностью контролировать эти аспекты своего обучения и не должны быть привязаны к библиотеке. Компромисс здесь заключается в том, что вам необходимо реализовать эти функции в своем обучающем сценарии (обычно это всего лишь две строки кода, чтобы уменьшить все ваши градиенты перед применением оптимизатора) или заменить модули, которые явно поддерживают распространение (например, snt.distribute.CrossReplicaBatchNorm ).
В нашем распределенном примере Cifar-10 показано обучение нескольких графических процессоров с помощью Sonnet.
