Многие люди, когда либо имевшие дело с нейронными сетями, наверняка задумывались, можно ли написать нейросеть, которая сама будет создавать нейросети для решения каких-либо задач. Так вот в этом цикле статей я решил реализовать это. Одним из этапов алгоритма будет генерирование нейросети из списка слоёв. В связи с некоторыми ограничениями, накладываемыми методами реализации (о которых будет сказано в следующих частях, когда мы начнём объединять код из этой статьи с RL ʕ⊙ᴥ⊙ʔ ), входные данные для генератора будут представлены в виде строки случайной длины, содержащей упорядоченный набор слоёв с их параметрами. Генерировать сеть будем для задачи классификации картинок (разобьём это пугало первым).

Функция создания описания последовательности слоёв

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

layers_set = {
                # increase channels with kernel 3
                'conv=channel_factor:2,kernel_size:3,stride:1,padding:0-',
                'conv=channel_factor:3,kernel_size:3,stride:1,padding:0-',
                'conv=channel_factor:4,kernel_size:3,stride:1,padding:0-',

                # decrease channels with kernel 3
                'conv=channel_factor:0.4,kernel_size:3,stride:1,padding:0-',
                'conv=channel_factor:0.6,kernel_size:3,stride:1,padding:0-',
                'conv=channel_factor:0.8,kernel_size:3,stride:1,padding:0-',


                 'batchnorm=eps:0.00001-',
                 'batchnorm=eps:0.0001-',

                 'avgpool=kernel_size:2,stride:2,padding:0-',
                 'avgpool=kernel_size:4,stride:4,padding:0-',

                 'maxpool=kernel_size:2,stride:2,padding:0-',
                 'maxpool=kernel_size:4,stride:4,padding:0-',

                 'dropout=p:0.2-',
                 'dropout=p:0.4-',
  }

channel_factor - множитель кол-ва каналов.
"-" - разделитель слоёв
"," - разделитель параметров слоя
":" - разделитель названия параметра и его значения
"=" - разделитель названия слоя от параметров

Реализуем функцию, создающую строку слоёв с длиной min_len <= L <= max_len слоёв.

def create_layers_string(min_len=2, max_len=10):
    if min_len < 1 or max_len < min_len:
      print('Parameters are incorrect!')
      return None
    length = random.randint(min_len, max_len)
    text = random.sample(layers_set, length)
    text = ''.join((layer for layer in text))
    return text

Класс нейросети

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

class NN(nn.Module):

    def __init__(self, in_channels):
      super().__init__()

      self.layers = nn.Sequential()

    def forward(self, x):
      x = self.layers(x)
      return x

    def __call__(self, x):
      return self.forward(x)

Замечу, что классификатора в виде линейного слоя тут нет. Он будет добавлен в конец layers после вычисления входных параметров для него.

Класс nnGenerator

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

class nnGenerator():
    def __init__(self):
      self.text_layers_dict = dict({})
      self.nn_len = -1

Класс содержит словарь, хранящий результат парсинга строки слоёв.
Структура словаря следующая: 1 элемент его состоит из название слоя, сконкатенированное с его id + пара: id + словарь параметров. Пример одного элемента словаря text_layers_dict:
"dropout1": (1, {"p": 0.2} )

    def parseTextNet(self, text_net):
      self.text_layers_dict = dict({})
      self.nn_len = -1
      print(text_net)
      if text_net[-1] == '-':
        text_net = text_net[:-1]
      text_layers = text_net.split('-')
      id = 0
      for text_layer in text_layers:
        tmp = text_layer.split('=')
        layer_name, layer_params = tmp[0], tmp[1].split(',')
        layer_params_dict = dict({})
        for param in layer_params:
          param = param.split(':')
          param_name, param_value = param[0], param[1]
          layer_params_dict[param_name] = param_value
        self.text_layers_dict[layer_name + str(id)] = (id, layer_params_dict)
        id += 1
      print(self.text_layers_dict)

Парсер принимает на вход строку слоёв и преобразует её, записывая результат в словарь text_layers_dict.

    def conv_output_shape(self, h, w, kernel_size=1, stride=1, pad=0, dilation=1):
      h = math.floor( ((h + (2 * pad) - ( dilation * (kernel_size - 1) ) - 1 )/ stride) + 1)
      w = math.floor( ((w + (2 * pad) - ( dilation * (kernel_size - 1) ) - 1 )/ stride) + 1)
      return h, w

Метод conv_output_shape реализует вычисление выхода нейросетевого слоя.
Для свёрток и пуллингов эта функция универсальна, dropout и batchnorm не меняют размерностей, а значит в вычислении выхода слоя не нуждаются. Формула вычисления этих значений представлена в документации библиотеки PyTorch.

Далее пойдёт самое интересное - автопостроение нашей сети по словарю слоёв.

    def generateNN(self, n_classes, test_batch):
      success_state = False

      backbone = nn.Sequential()
      classifier = nn.Sequential()
      optimizer = None
      try:
        data_shape = np.array(test_batch).shape # [B, C, H, W]
        last_shape = data_shape
        for layer_name in self.text_layers_dict.keys():
            layer = None
            layer_params= self.text_layers_dict[layer_name][1]
            if layer_name.find('conv') >= 0:
              kernel_size = int(layer_params['kernel_size'])
              channel_factor = float(layer_params['channel_factor'])
              stride = int(layer_params['stride'])
              padding = int(layer_params['padding'])
              dilation = 1

              activation = nn.ReLU(inplace=True)

              in_chan = last_shape[1]
              assert(in_chan <= last_shape[2] and in_chan <= last_shape[3])
              out_chan = math.floor(in_chan * channel_factor)
              assert(out_chan > 0)

              backbone.append(nn.Conv2d(in_chan, out_chan, kernel_size, stride, padding, dilation))
              backbone.append(activation)

              h, w = self.conv_output_shape(last_shape[2], last_shape[3], kernel_size, stride, padding, dilation)
              last_shape = (last_shape[0], out_chan, h, w)

            elif layer_name.find('batchnorm') >= 0:
              eps = float(layer_params['eps'])
              in_chan = last_shape[1]
              backbone.append(nn.BatchNorm2d(in_chan, eps))

            elif layer_name.find('avgpool') >= 0:
              kernel_size = int(layer_params['kernel_size'])
              stride = int(layer_params['stride'])
              padding = int(layer_params['padding'])
              dilation = 1
              out_chan = last_shape[1]
              backbone.append(nn.AvgPool2d(kernel_size, stride, padding))

              h, w = self.conv_output_shape(last_shape[2], last_shape[3], kernel_size, stride, padding, dilation)
              last_shape = (last_shape[0], out_chan, h, w)


            elif layer_name.find('maxpool') >= 0:
              kernel_size = int(layer_params['kernel_size'])
              stride = int(layer_params['stride'])
              padding = int(layer_params['padding'])
              dilation = 1

              backbone.append(nn.MaxPool2d(kernel_size, stride, padding))

              h, w = self.conv_output_shape(last_shape[2], last_shape[3], kernel_size, stride, padding, dilation)
              last_shape = (last_shape[0], last_shape[1], h, w)

            elif layer_name.find('dropout') >= 0:
              p = float(layer_params['p'])
              backbone.append(nn.Dropout2d(p))

        linear_in_shape = last_shape[1] * last_shape[2] * last_shape[3]
        classifier = nn.Linear(linear_in_shape, n_classes)
        success_state = True
        print('NN build successfull!')
        
      except Exception as e:
        print('NN build failed!')
        print(str(e))

      net = nn.Sequential()
      net.append(backbone)
      net.append(nn.Flatten(start_dim=1))
      net.append(classifier)
      self.text_layers_dict = dict({})
      return success_state, net

Псевдокодом можно описать алгоритм так:
Проходим по всем слоям из словаря:
1) Берём параметры слоя
2) Вычисляем выходной shape
3) Создаём слой, добавляем в конец архитектуры
4) Сохраняем значение выходного размера последнего слоя
В конце добавляем линейный классификатор.
Для вычисления входного размера данных подаётся тестовый батч, на ваше усмотрение можно переделать в более удобный вид.

Особые моменты, которые необходимо отлавливать, это случай, при котором в ходе построения сети ядро становится больше размера данных, и при котором кол-во каналов уменьшается до значения < 1.

Далее код создания dataloader'ов под MNIST, а также обучающего цикла, думаю, в комментариях это не нуждается, отмечу только, что входные данные должны быть (B, C, H, W) = (Batch_Size, Channels, height, width), значит данные из МНИСТа с shape'ом (N, 28,28) нужно представить в виде (N, 1, 28, 28)

class ds(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_

train_data = dsets.MNIST(root = './data', train = True,
                        transform = transforms.ToTensor(), download = True)

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


train_samples = np.expand_dims(np.array(train_data.data), axis=1)[:5000]
train_labels = np.array(train_data.targets)[:5000]


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

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

train_dataloader = torch.utils.data.DataLoader(
  train_set,
  batch_size=bs,
  shuffle=True,
  drop_last=True)

valid_dataloader = torch.utils.data.DataLoader(
  valid_set,
  batch_size=bs,
  drop_last=True,
  shuffle=True)

test_batch = None
for b, _ in train_dataloader:
  test_batch = b
  break

Тесты

Генерируем последовательность слоёв

test_batch = None
for b, _ in train_dataloader:
  test_batch = b
  break

generator = nnGenerator()

success_state = False
while not success_state:
  text_layers = create_layers_string()
  generator.parseTextNet(text_layers)
  success_state, seq = generator.generateNN(n_classes=10, test_batch=test_batch)

Результатом будет текст следующего содержания:

conv=channel_factor:0.4,kernel_size:3,stride:1,padding:0-batchnorm=eps:0.00001-conv=channel_factor:4,kernel_size:3,stride:1,padding:0-conv=channel_factor:0.8,kernel_size:3,stride:1,padding:0-maxpool=kernel_size:4,stride:4,padding:0-
{'conv0': (0, {'channel_factor': '0.4', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'batchnorm1': (1, {'eps': '0.00001'}), 'conv2': (2, {'channel_factor': '4', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'conv3': (3, {'channel_factor': '0.8', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'maxpool4': (4, {'kernel_size': '4', 'stride': '4', 'padding': '0'})}
NN build failed!

batchnorm=eps:0.00001-conv=channel_factor:3,kernel_size:3,stride:1,padding:0-avgpool=kernel_size:4,stride:4,padding:0-conv=channel_factor:0.6,kernel_size:3,stride:1,padding:0-maxpool=kernel_size:4,stride:4,padding:0-dropout=p:0.2-conv=channel_factor:0.8,kernel_size:3,stride:1,padding:0-dropout=p:0.4-
{'batchnorm0': (0, {'eps': '0.00001'}), 'conv1': (1, {'channel_factor': '3', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'avgpool2': (2, {'kernel_size': '4', 'stride': '4', 'padding': '0'}), 'conv3': (3, {'channel_factor': '0.6', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'maxpool4': (4, {'kernel_size': '4', 'stride': '4', 'padding': '0'}), 'dropout5': (5, {'p': '0.2'}), 'conv6': (6, {'channel_factor': '0.8', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'dropout7': (7, {'p': '0.4'})}
NN build failed!

avgpool=kernel_size:4,stride:4,padding:0-conv=channel_factor:0.6,kernel_size:3,stride:1,padding:0-dropout=p:0.4-conv=channel_factor:0.8,kernel_size:3,stride:1,padding:0-
{'avgpool0': (0, {'kernel_size': '4', 'stride': '4', 'padding': '0'}), 'conv1': (1, {'channel_factor': '0.6', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'dropout2': (2, {'p': '0.4'}), 'conv3': (3, {'channel_factor': '0.8', 'kernel_size': '3', 'stride': '1', 'padding': '0'})}
NN build failed!

conv=channel_factor:0.6,kernel_size:3,stride:1,padding:0-avgpool=kernel_size:2,stride:2,padding:0-conv=channel_factor:0.8,kernel_size:3,stride:1,padding:0-conv=channel_factor:0.4,kernel_size:3,stride:1,padding:0-conv=channel_factor:4,kernel_size:3,stride:1,padding:0-conv=channel_factor:3,kernel_size:3,stride:1,padding:0-maxpool=kernel_size:4,stride:4,padding:0-batchnorm=eps:0.00001-maxpool=kernel_size:2,stride:2,padding:0-batchnorm=eps:0.0001-
{'conv0': (0, {'channel_factor': '0.6', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'avgpool1': (1, {'kernel_size': '2', 'stride': '2', 'padding': '0'}), 'conv2': (2, {'channel_factor': '0.8', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'conv3': (3, {'channel_factor': '0.4', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'conv4': (4, {'channel_factor': '4', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'conv5': (5, {'channel_factor': '3', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'maxpool6': (6, {'kernel_size': '4', 'stride': '4', 'padding': '0'}), 'batchnorm7': (7, {'eps': '0.00001'}), 'maxpool8': (8, {'kernel_size': '2', 'stride': '2', 'padding': '0'}), 'batchnorm9': (9, {'eps': '0.0001'})}
NN build failed!

conv=channel_factor:0.6,kernel_size:3,stride:1,padding:0-dropout=p:0.2-batchnorm=eps:0.00001-maxpool=kernel_size:4,stride:4,padding:0-avgpool=kernel_size:2,stride:2,padding:0-conv=channel_factor:0.4,kernel_size:3,stride:1,padding:0-conv=channel_factor:4,kernel_size:3,stride:1,padding:0-avgpool=kernel_size:4,stride:4,padding:0-conv=channel_factor:0.8,kernel_size:3,stride:1,padding:0-batchnorm=eps:0.0001-
{'conv0': (0, {'channel_factor': '0.6', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'dropout1': (1, {'p': '0.2'}), 'batchnorm2': (2, {'eps': '0.00001'}), 'maxpool3': (3, {'kernel_size': '4', 'stride': '4', 'padding': '0'}), 'avgpool4': (4, {'kernel_size': '2', 'stride': '2', 'padding': '0'}), 'conv5': (5, {'channel_factor': '0.4', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'conv6': (6, {'channel_factor': '4', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'avgpool7': (7, {'kernel_size': '4', 'stride': '4', 'padding': '0'}), 'conv8': (8, {'channel_factor': '0.8', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'batchnorm9': (9, {'eps': '0.0001'})}
NN build failed!

batchnorm=eps:0.00001-dropout=p:0.4-batchnorm=eps:0.0001-conv=channel_factor:0.8,kernel_size:3,stride:1,padding:0-maxpool=kernel_size:4,stride:4,padding:0-conv=channel_factor:3,kernel_size:3,stride:1,padding:0-
{'batchnorm0': (0, {'eps': '0.00001'}), 'dropout1': (1, {'p': '0.4'}), 'batchnorm2': (2, {'eps': '0.0001'}), 'conv3': (3, {'channel_factor': '0.8', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'maxpool4': (4, {'kernel_size': '4', 'stride': '4', 'padding': '0'}), 'conv5': (5, {'channel_factor': '3', 'kernel_size': '3', 'stride': '1', 'padding': '0'})}
NN build failed!

dropout=p:0.2-conv=channel_factor:3,kernel_size:3,stride:1,padding:0-dropout=p:0.4-
{'dropout0': (0, {'p': '0.2'}), 'conv1': (1, {'channel_factor': '3', 'kernel_size': '3', 'stride': '1', 'padding': '0'}), 'dropout2': (2, {'p': '0.4'})}
NN build successfull!
<ipython-input-20-4e57f99010fc>:6: DeprecationWarning: Sampling from a set deprecated
since Python 3.9 and will be removed in a subsequent version.
  text = random.sample(layers_set, length)

С седьмой попытки получилось (；一_一)

Создаём объект пустой сети, генерируем сеть

model = NN(in_channels=1)
model.layers = seq
display(model)

Получаем какую-то нашу случайную сеть

NN(
  (layers): Sequential(
    (0): Sequential(
      (0): Dropout2d(p=0.2, inplace=False)
      (1): Conv2d(1, 3, kernel_size=(3, 3), stride=(1, 1))
      (2): ReLU(inplace=True)
      (3): Dropout2d(p=0.4, inplace=False)
    )
    (1): Flatten(start_dim=1, end_dim=-1)
    (2): Linear(in_features=2028, out_features=10, bias=True)
  )
)

Далее оптимайзер + ф-я ошибки

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

Тренировочный цикл самый простой

def train():
    train_losses = []
    valid_losses = []
    # TODO calculate metrics and return them after train
    def CalcValLoss():
        with torch.no_grad():
            losses = []
            for X, Y in valid_dataloader:
                X = X.float().to(device)
                Y = Y.float().to(device)
                preds = model(X)
                preds, _ = torch.max(preds,1)
                loss = criterion(preds,Y)
                losses.append(loss.item())
            print("Valid Loss : {:.6f}".format(torch.tensor(losses).mean()))
            valid_losses.append(torch.tensor(losses).mean())

    for i in range(1, epochs):
        losses = []
        for X, Y in tqdm(train_dataloader):
            X = X.float().to(device)
            Y = Y.float().to(device)
            preds = model(X)
            preds, _ = torch.max(preds,1)
            loss = criterion(preds, Y)
            losses.append(loss.item())

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print("Train Loss : {:.6f}".format(torch.tensor(losses).mean()))
        train_losses.append(torch.tensor(losses).mean())
    return train_losses, valid_losses

Запускаем, обучаем, фиксируем прибыль

train_losses, valid_losses = train()

100%|██████████| 125/125 [00:00<00:00, 213.39it/s]
Train Loss : 8579.939453
100%|██████████| 125/125 [00:00<00:00, 219.21it/s]
Train Loss : 493.860870
100%|██████████| 125/125 [00:00<00:00, 224.44it/s]
Train Loss : 493.228210
100%|██████████| 125/125 [00:00<00:00, 347.04it/s]
Train Loss : 493.066498
100%|██████████| 125/125 [00:00<00:00, 359.57it/s]
Train Loss : 492.938690
100%|██████████| 125/125 [00:00<00:00, 340.28it/s]
Train Loss : 492.947876
100%|██████████| 125/125 [00:00<00:00, 339.37it/s]
Train Loss : 492.833740
100%|██████████| 125/125 [00:00<00:00, 335.57it/s]
Train Loss : 492.833679
100%|██████████| 125/125 [00:00<00:00, 350.72it/s]
Train Loss : 492.908295

Итог

Мы получили генератор свёрточных нейронных сетей случайной длины для классификации изображений.

Планы

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

Ссылка на гитхаб с ipynb файлом кода.

Тэги: pytorch, RL, reinforcement learning, генерация нейронных сетей, нейронная сеть, neural network.