Нейросеть, генерирующая нейросети. Часть 2. RL агент создаёт свои первые нейросети

Введение

Наконец пришло время объединить код из предыдущей части, в которой мы создавали нейросеть по списку слоёв, с 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  ')         for name in stage_names:             register_str = "register(     id='{}',     entry_point='gym_{}.envs:{}', ) ".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 {} ".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  ',             "setup(name='gym_{}', ".format(env_name),             "    version='0.0.1', ",             "    install_requires=['gym']  # And any other dependencies foo needs ",             ") ",         ]         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 (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