Введение

Наконец пришло время объединить код из предыдущей части, в которой мы создавали нейросеть по списку слоёв, с RL алгоритмами! Сегодня мы поставим задачу для обучения с подкреплением, опишем, как будет производиться взаимодействие агента со средой и на практике реализуем код на языке python.

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

Начнём с постановки задачи

Первым делом опишем, что мы хотим получить в результате.

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

Упрощения для первой реализации

Так задача имеет много параметров, от которых зависит качество результата, необходимо их определить и некоторые зафиксировать.

Гиперпараметры для обучения нейросети

learning_rate=0.001
batch_size=32
input_shape=[B, C, H, W]
num_classes=10 (датасет CIFAR10)
train_epochs=10
device='cuda' (Тестировалось на T4 GPU в google colab)

Гиперпараметры среды

layers_amount=5 - длина последовательности слоёв, генерируемой агентом.

memory_len=10 - сколько последних наблюдений будут выдаваться агенту, длина кольцевого буфера, новые измерения добавляются в конец.

metrics_amount=2 (Будем использовать просто train и valid MSELosses, да, не оптимально, но это в первой реализации не принципиальный момент).

Observation, action, reward

В качестве входных данных для агента (observation) мы возьмём метрики, полученные в ходе обучения, и архитектуры в количестве, равном длине памяти. Т.е. агент принимает решение о построении новых нейросетей на основе пар: (предыдущая архитектура, метрики после обучения)

obs={
'metrics': shape=[mem_len, amount_of_metrics, train_epochs],
'architectures': shape=[mem_lem, layers_amount],
}

Действие агента (action) будет выглядеть как набор чисел фиксированной длины, каждому из которых соответствует слой нейросети из списка возможных (actions_set).

action = []
action.shape=[max_layers_amount]

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

• metrics_optimization_reward - в случае, когда метрики улучшаются, агент заслуживает поощрение. рассчитаем награду как разницу между последними значениями отслеживаемых метрик и новыми, в зависимости от знака получившегося числа награда может превратиться в наказание за построение модели хуже, чем в последний раз.
  Замечание 1: метрики первой эпохи могут представлять собой выбросы большой величины, поэтому их мы будем пропускать, производя по сути train_epochs + 1 эпох обучения.
  Замечание 2: Не самые оптимальные архитектуры могут выдавать огромные показатели метрики. Для сглаживания этого момента мы будем ограничивать сверху значения эмпирически подобранным порогом (в частности было использовано значение 100).
  Замечание 3: иногда в особо печальных случаях метрики вообще могут принять значение nan или inf, будем также заменять такие значения на NAN_NUM = 100.

• creation_successfull_reward - если нейросеть вообще получилось создать, уже неплохо, будем награждать агента за нейросети, которые получилось корректно построить и штрафовать в обратном случае. Для этого были созданы переменные NN_CREATE_SUCCESS_REWARD = 1
  NN_CREATE_NOT_SUCCESS_PENALTY = -1

• optimal_depth_reward - эта часть является заделом на будущее, на текущий момент не используется

RL Custom environment

Для реализации будет использоваться python с библиотеками gymnasium и stable-baselines3.

Напомню основную структуру среды в gymnasium:

class Env(gym.Env):
  def __init__(self):
    # init all required variables
    # action and observation spaces must be here!

    self.action_space = spaces.Box(...)
    self.observation_space = spaces.MultiDiscrete(...)
    pass

  def seed(self):
     seed = 42
     return [seed]

  def step(self, action):
    # most important metrics
    # do manipulations with action, 
    # return new observation and reward
    return obs, reward, _, done, info

  def reset(self):
    # set variables to start state
    return obs

  def render(self):
    # all GUI parts should be here
    # may be empty
    pass

Пройдёмся по каждому пункту.

__init__, объявление переменных

  def __init__(self, opt_cls=None, crit=None, trn_ldr=None, vld_ldr=None, render_mode=None):
      '''
      Initialization of environment variables
      '''

      # here is learning params
      self.NN_PARAMS = {
          'lr': 0.001,
          'num_classes': 10, # TODO adaptive num_classes
          'train_epochs': 10,
          'last_nets_metrics_memory_len': 10,
          'layers_amount': 5,
          'amount_of_metrics': 2,
      }

      # here contains NN learning metrics
      # uses for observations
      # shape=[MEMORY_LEN, N_METRICS, EPOCHS]
      self.NN_PARAMS['metrics'] = np.zeros((
          self.NN_PARAMS['last_nets_metrics_memory_len'],
          self.NN_PARAMS['amount_of_metrics'],
          self.NN_PARAMS['train_epochs'],
      ))

      # here contains last architectures
      # uses for observations
      # shape=[MEMORY, N_LAYERS]
      self.NN_PARAMS['last_nets_architectures'] = np.zeros((
             self.NN_PARAMS['last_nets_metrics_memory_len'],
             self.NN_PARAMS['layers_amount'],
          ))

      # Variables for NN
      self.Net = self.NN()
      self.train_dataloader = trn_ldr
      self.valid_dataloader = vld_ldr
      self.optimizer_class = opt_cls
      self.optimizer = None
      self.criterion = crit
      self.device = 'cuda'

      self.nngenerator = self.nnGenerator()

      self.last_obs = None

      # Action space describes what agent will give to environment
      # shape=[ACTION_SET_SIZE, N_LAYERS]
      self.action_space = spaces.MultiDiscrete(
        [len(self.actions_set)] * self.NN_PARAMS['layers_amount'],
         seed=42)

      # Observation space describes what agent
      # will take from enviromnent as observation
      # shape=dict{
      #  METRICS, shape as NN_PARAMS['metrics'],
      #  ARCHITECTURES, shape as NN_PARAMS['layers_amount'],
      # }
      self.observation_space = spaces.Dict(
          {
          'last_nets_metrics_memory': spaces.Box(
              low=0,
              high=100,
              shape=(self.NN_PARAMS['last_nets_metrics_memory_len'],
                     self.NN_PARAMS['amount_of_metrics'],
                     self.NN_PARAMS['train_epochs'],
                     )),
          'last_nets_architectures': spaces.Box(
              low=0,
              high=len(self.actions_set),
              shape=(self.NN_PARAMS['last_nets_metrics_memory_len'],
                     self.NN_PARAMS['layers_amount'],
                     ))
          }
      )

      # some variable for collecting statistics

      self.statistics = {
        'episode_rewards': [],
        'global_rewards': [],
        'made_steps': [],
      }

      self.seed()
      assert render_mode is None or render_mode in self.metadata["render_modes"]
      self.render_mode = render_mode

На вход принимаются оптимайзер, ф‑я ошибки и даталоадеры.

Здесь указаны все необходимые параметры, создаётся объект пустой нейросети Net, в которую мы позже будет помещать построенную нейросеть, а также nnGenerator'a, что занимается обработкой действия агента и построения самой сети. Обязательными полями являются action_space и observation_space. Без их определения создание среды невозможно.

Замечание: к выбору пространства действий и наблюдений следует подходить с умом, так как алгоритмы агента из stable‑baselines3 могут не поддерживать тот формат, который вы выберете. О том, какие пространства можно выбрать для алгоритмов, можно посмотреть тут.
Мы будем использовать PPO('MultiInputPolicy') — proximal policy optimization алгоритм, о котором можно прочитать в документации тут.

reset, обнуление переменных, перезагрузка среды

  def reset(self, seed=None, options=None):
      '''
      Reset the env,
      Set all changed in training proccess variables to zero
      (or noise, dependse on your realization)

      seed: list, list of random seeds (depricated)
      options: list, additional options (future)
      '''
      super().reset(seed=seed)
      self.Net = self.NN()

      current_obs = {
          'last_nets_metrics_memory': np.zeros((
              self.NN_PARAMS['last_nets_metrics_memory_len'],
              self.NN_PARAMS['amount_of_metrics'],
              self.NN_PARAMS['train_epochs'],
             ) ),
          'last_nets_architectures': np.zeros((
             self.NN_PARAMS['last_nets_metrics_memory_len'],
             self.NN_PARAMS['layers_amount'],
          ) )

          }
      self.NN_PARAMS['metrics'] = np.zeros((
          self.NN_PARAMS['last_nets_metrics_memory_len'],
          self.NN_PARAMS['amount_of_metrics'],
          self.NN_PARAMS['train_epochs'],
      ))
      self.NN_PARAMS['last_nets_architectures'] = np.zeros((
             self.NN_PARAMS['last_nets_metrics_memory_len'],
             self.NN_PARAMS['layers_amount'],
          ))

      self.last_obs = current_obs
      self.episode_reward = 0
      self.current_it = 1

      self.statistics['episode_rewards'] = []
      self.statistics['made_steps'] = []

      return current_obs, {'none': None}

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

step, переводим action в new_obs и reward

  def step(self, action):
      '''
          Main environment function

          takes action, creates NN, train NN, calc new obs and reward

          action: list, shape of action is like action_space

      return:
      obs: np.array,
      reward: float,
      done = False, end of agent training flag, needed when your actions
      achieved some finish state
      info: dict, you may need to add some extra information, put it here

      Algorithm
      1) parse action
      2) generate NN
      3) update optimizer, prepare for training
      4) train NN, collect metrics
      5) calculate reward
      6) collect statistics
      7) create new observation
      8) return obs, reward, done, info

      '''
      reward = 0
      done = False
      info = {}
      self.nngenerator.parse_action(action, self.actions_set)
      test_b = self.get_test_batch()
      success_state, net = self.nngenerator.generateNN(n_classes=self.NN_PARAMS['num_classes'], test_batch=test_b)

      new_metrics = None
      if success_state == True: # NN created_correctly
        self.NN_PARAMS['last_nets_architectures'] = np.roll(self.NN_PARAMS['last_nets_architectures'], -1, axis=0)
        self.NN_PARAMS['last_nets_architectures'][-1] = action
        self.Net.layers = net
        display(self.Net)
        self.Net = self.Net.to(self.device)
        self.optimizer = self.optimizer_class(self.Net.parameters(), lr=self.NN_PARAMS['lr'])
        new_metrics = self.train()


      reward = self.calc_reward(success_state,
                                self.nngenerator.get_nn_len(),
                                new_metrics,
                                )


      current_obs = self.create_obs()

      self.last_obs = current_obs

      print('Reward: ', reward)
      self.episode_reward = reward
      self.current_it += 1
      return current_obs, self.episode_reward, None, done, info

В этом методе происходит вся магия. Создаём нейросеть, обучаем, собираем метрики, считаем награду, фиксируем прибыль ヽ( ▀̿ Ĺ̯ ▀̿)ノ.

Вспомогательными методами здесь являются get_test_batch, calc_reward и create_obs.
Выглядят они следующим образом:

  def calc_reward(self, nn_created_correctly_flag, nn_len, last_train_metrics):
    '''
       calculate agent reward

       nn_created_correctly_flag: bool,
       if True - NN was built successfully,
       we can calculate other parts of reward,
       otherwise - agent takes NN_CREATE_NOT_SUCCESS_PENALTY only

       nn_len: int,
       that var needed for depth decreasing reward
       # Not used (future) #

       last_train_metrcis: np.array,
       contains last training loop metrics
       for metrics_optimization_reward

       return reward: float, sum of all reward parts

    '''

    reward = 0
    # TODO reward for decreasing nn depth
    optimal_depth_reward = 0
    # reward by metrics
    metrics_optimization_reward = 0
    # reward for successfull nn creation
    creation_successfull_reward = 0
    if nn_created_correctly_flag == True:
      # do not reward agent if creation is not succeed
      creation_successfull_reward += self.NN_CREATE_SUCCESS_REWARD

      last_metrics = self.NN_PARAMS['metrics'][-1]
      last_train_metrics = np.nan_to_num(x=last_train_metrics, nan=self.NAN_NUM)
      last_train_metrics = np.minimum(last_train_metrics,
                          np.ones(shape=(self.NN_PARAMS['amount_of_metrics'], self.NN_PARAMS['train_epochs'])) * self.NAN_NUM
                          )

      tmp_r = np.min(last_metrics,axis=1) - np.min(np.array(last_train_metrics), axis=1)


      self.NN_PARAMS['metrics'] = np.roll(self.NN_PARAMS['metrics'], -1, axis=0)
      self.NN_PARAMS['metrics'][-1] = last_train_metrics
      metrics_optimization_reward = np.sum(self.METRICS_OPTIMIZATION_FACTOR * tmp_r)

      optimal_depth_reward = (self.NN_PARAMS['layers_amount'] - nn_len)
      optimal_depth_reward *= self.DEPTH_REWARD_FACTOR

    else:
      creation_successfull_reward += self.NN_CREATE_NOT_SUCCESS_PENALTY

    reward += optimal_depth_reward
    reward += metrics_optimization_reward
    reward += creation_successfull_reward

    return reward



  def get_test_batch(self):
    '''
        return: batch: np.array, one batch
        for input_shape in NN building algorithm

    '''
    batch = None
    for  b, _ in self.train_dataloader:
      batch = b
    return batch




  def create_obs(self):
    '''
       return obs: dict, new observation
    '''

    obs = {
          'last_nets_metrics_memory': self.NN_PARAMS['metrics'],
          'last_nets_architectures': self.NN_PARAMS['last_nets_architectures']

          }
    return obs

get_test_batch достаёт из датасета один батч, необходимый для вычисления входного shape'a данных при построении нейросети.

create_obs возвращает новое наблюдение, объединяя метрики и архитектуры в словарь формата, заданного в observation_space.

calc_reward посчитывает, какой награды (или наказания) будет удостоен агент.

Тестирование

Первым делом нужно зарегистрировать нашу самодельную среду в gymnasium библиотеке.
Для этого мы обернём наш код среды особым образом, чтобы все необходимые файлы сгенерировались. Вот так это выглядит с шаблоном gym среды, описанном выше:

import os


def make_dirs(env_name: str):
    os.makedirs('gym_{}'.format(env_name))
    os.makedirs('gym_{}/envs'.format(env_name))
    pass


def make_files(env_name: str, stage_names):
    registry_init = 'gym_{}/__init__.py'.format(env_name)
    with open(registry_init, 'w') as f:
        f.write('from gym.envs.registration import register\n\n')
        for name in stage_names:
            register_str = "register(\n    id='{}',\n    entry_point='gym_{}.envs:{}',\n)\n".format(name, 
                                                                                                    env_name,
                                                                                                    ''.join([s.capitalize()for s in name.split('_')]))
            f.write(register_str)
    envs_init = 'gym_{}/envs/__init__.py'.format(env_name)
    with open(envs_init, 'w') as f:
        for name in stage_names:
            import_str = "from gym_{}.envs.{} import {}\n".format(env_name, name,
                                                                  ''.join([s.capitalize() for s in name.split('_')]))
            f.write(import_str)

    head = """#all required imports here

"""

    tail = """
  def __init__(self):
    # init all required variables
    # action and observation spaces must be here!

    self.action_space = spaces.Box(...)
    self.observation_space = spaces.MultiDiscrete(...)
    pass

  def seed(self):
     seed = 42
     return [seed]

  def step(self, action):
    # most important metrics
    # do manipulations with action, 
    # return new observation and reward
    return obs, reward, _, done, info

  def reset(self):
    # set variables to start state
    return obs

  def render(self):
    # all GUI parts should be here
    # may be empty
    pass
  

    """
    for name in stage_names:
        env_class = 'class {}(gym.Env):'.format(''.join([s.capitalize() for s in name.split('_')]))
        with open('gym_{}/envs/{}.py'.format(env_name, name), 'w') as f:
            f.write(head + env_class + tail)
    pass


def make_setup(env_name):
    with open('setup.py', 'w') as f:
        tmp = [
            'from setuptools import setup\n\n',
            "setup(name='gym_{}',\n".format(env_name),
            "    version='0.0.1',\n",
            "    install_requires=['gym']  # And any other dependencies foo needs\n",
            ")\n",
        ]
        f.writelines(tmp)


if __name__ == '__main__':
    env_name = 'env'# input('Please type your env\'s name (string):')
    env_stage_num = 1 #int(input('Type the stage number of your Env (int):'))

    print(
        'Please type yout stage names\n(start with alphabet, lowercase and number, split by underscore(_), e.g. my_env_v1)')
    env_stage_names = ['env']
    #for i in range(env_stage_num):
    #    env_stage_names.append(input("Stage {}s name:".format(i)))

    make_dirs(env_name)
    make_files(env_name, env_stage_names)
    #generate_setup = input('Generate `setup.py`?[y/n]')

    generate_setup = 'y'
    if generate_setup.lower() == 'y':
        make_setup(env_name)

Далее мы регистрируем нашу среду в gym, создаём её и запускаем обучение:

# for google colab
#!pip install stable_baselines3
#!pip install shimmy==0.2.1
#!pip install gymnasium
from stable_baselines3 import *
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
from stable_baselines3.common.env_util import make_vec_env
import gymnasium as gym
import numpy as np
import torch
import torchvision.datasets as dsets
from torchvision import transforms
from torch.utils.data import Dataset
from gymnasium.envs.registration import register

register(
    id='gym_env/env_v1-v0',
    entry_point='gym_env.envs:Env',
    max_episode_steps=300,
)

# download pytorch CIFAR10 Dataset
train_data = dsets.CIFAR10(root = './data', train = True,
                        transform = transforms.ToTensor(), download = True)

test_data = dsets.CIFAR10(root = './data', train = False,
                       transform = transforms.ToTensor())

train_samples = np.array(train_data.data)[:15000].transpose((0,3,1,2)) # convert to [B, C, H, W]
train_labels = np.array(train_data.targets)[:15000]

# simple custom dataset template
class myDataset(Dataset):
    def __init__(self, X, y):

      self.X = X
      self.y = y

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        x_ = self.X[idx]
        y_ = self.y[idx]
        return x_, y_

dataset = myDataset(X=train_samples, y=train_labels)

train_set, valid_set = torch.utils.data.random_split(dataset, [0.8, 0.2], generator=torch.Generator().manual_seed(42))


train_loader = torch.utils.data.DataLoader(
  train_set,
  batch_size=32,
  shuffle=True,
  drop_last=True)

valid_loader = torch.utils.data.DataLoader(
  valid_set,
  batch_size=32,
  drop_last=True,
  shuffle=True)


# some environment params
# optimizer should be class
# because every NN building needs
# pass model.params() to optimizer,
# so we will recreate optimizer every
# time new NN will be created successfully

env_params = {
    'opt_cls': torch.optim.RAdam,
    'trn_ldr': train_loader,
    'vld_ldr': valid_loader,
    'crit': torch.nn.MSELoss(),
}
env = gymnasium.make('gym_env/env_v1', **env_params)
# Proximal Policy Optimization
# doc: https://stable-baselines3.readthedocs.io/en/master/modules/ppo.html

model = PPO("MultiInputPolicy", env, verbose=1)

model.learn(total_timesteps=10000)

Процесс обучения выглядит следующим образом:

Reward:  -1
Reward:  -1
Reward:  -1
Reward:  -1
Reward:  -1
NN(
  (layers): Sequential(
    (0): Sequential(
      (0): AvgPool2d(kernel_size=2, stride=2, padding=0)
      (1): Dropout2d(p=0.1, inplace=False)
      (2): Dropout2d(p=0.1, inplace=False)
      (3): Conv2d(3, 2, kernel_size=(5, 5), stride=(1, 1))
      (4): ReLU(inplace=True)
      (5): Conv2d(2, 16, kernel_size=(5, 5), stride=(1, 1))
      (6): ReLU(inplace=True)
    )
    (1): Flatten(start_dim=1, end_dim=-1)
    (2): Linear(in_features=1024, out_features=10, bias=True)
  )
)
100%|██████████| 11/11 [00:10<00:00,  1.04it/s]Reward:  -3.344358348846436

NN(
  (layers): Sequential(
    (0): Sequential(
      (0): Conv2d(3, 2, kernel_size=(5, 5), stride=(1, 1))
      (1): ReLU(inplace=True)
      (2): Dropout2d(p=0.1, inplace=False)
      (3): Dropout2d(p=0.1, inplace=False)
      (4): Dropout2d(p=0.2, inplace=False)
      (5): AvgPool2d(kernel_size=2, stride=2, padding=0)
    )
    (1): Flatten(start_dim=1, end_dim=-1)
    (2): Linear(in_features=392, out_features=10, bias=True)
  )
)
100%|██████████| 11/11 [00:09<00:00,  1.12it/s]
Reward:  -5.165955924987793
Reward:  -1
Reward:  -1
Reward:  -1
Reward:  -1
Reward:  -1
Reward:  -1
NN(
  (layers): Sequential(
    (0): Sequential(
      (0): Conv2d(3, 1, kernel_size=(5, 5), stride=(1, 1))
      (1): ReLU(inplace=True)
      (2): Dropout2d(p=0.2, inplace=False)
      (3): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (4): Conv2d(1, 8, kernel_size=(5, 5), stride=(1, 1))
      (5): ReLU(inplace=True)
      (6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (1): Flatten(start_dim=1, end_dim=-1)
    (2): Linear(in_features=1152, out_features=10, bias=True)
  )
)
100%|██████████| 11/11 [00:11<00:00,  1.06s/it]
Reward:  10.80641279220581

Задел на будущее

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

Заключение

Обучение RL моделей — процесс довольно долгий, положительный результат может прийти далеко не сразу. Однако в перспективе получить нейросеть, которая сама будет решать задачу построения нейросети, вполне реально. Кто знает, возможно скоро вместо ML инженеров нейросети будут строить сами нейросети?

Репо с файлом исходного кода в формате ipynb