AI-движки на примере Knowledge Distillation, GAN, Reinforcement learning

Привет хабр!

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

Возьму как пример 3 подхода :

1. Архитектура GAN, основанная на состязательности нейросетей

2. Архитектура Knowledge Distillation, основанная на совместном обучении и дистилляции

3. Архитектура Reinforcement learning, основанная на последовательной или разделенной обработке

1. GAN – Генеративно – состязательные сети.

Генератор – это сеть, что получает на вход, так называемые, скрытые переменные (latent space) (случайный шум), а на выходе получаются данные ( изображение). Проще говоря постоянно рисует новые картинки, стараясь сделать их максимально похожими на настоящие.

Подгружаем датасет MNIST и смотрим случайную картинку из него, на и нем будем отрабатывать обучение

(X_train, y_train), (_, _) = tf.keras.datasets.mnist.load_data()
i = np.random.randint(0, 60000)
print(y_train[i])
plt.imshow(X_train[i], cmap='gray');

Смотрим предварительно случайную картинку из всего датасета

Cоздаем сеть Генератора

def build_generator():
  network = tf.keras.Sequential()

  network.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100, )))
  network.add(layers.BatchNormalization())
  network.add(layers.LeakyReLU())

  network.add(layers.Reshape((7, 7, 256)))

  # 7x7x128
  network.add(layers.Conv2DTranspose(128, (5,5), padding='same', use_bias=False))
  network.add(layers.BatchNormalization())
  network.add(layers.LeakyReLU())

  # 14x14x64
  network.add(layers.Conv2DTranspose(64, (5,5), strides = (2,2), padding='same', use_bias=False))
  network.add(layers.BatchNormalization())
  network.add(layers.LeakyReLU())

  # 28x28x1
  network.add(layers.Conv2DTranspose(1, (5,5), strides = (2,2), padding='same', use_bias=False, activation='tanh'))

  network.summary()

  return network

Билдим и смотрим таблицу, как сформировались слои Генератора

generator = build_generator()

Model: "sequential_1"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ dense_1 (Dense)                 │ (None, 12544)          │     1,254,400 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ batch_normalization_3           │ (None, 12544)          │        50,176 │
│ (BatchNormalization)            │                        │               │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ leaky_re_lu_3 (LeakyReLU)       │ (None, 12544)          │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ reshape_1 (Reshape)             │ (None, 7, 7, 256)      │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ conv2d_transpose_3              │ (None, 7, 7, 128)      │       819,200 │
│ (Conv2DTranspose)               │                        │               │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ batch_normalization_4           │ (None, 7, 7, 128)      │           512 │
│ (BatchNormalization)            │                        │               │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ leaky_re_lu_4 (LeakyReLU)       │ (None, 7, 7, 128)      │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ conv2d_transpose_4              │ (None, 14, 14, 64)     │       204,800 │
│ (Conv2DTranspose)               │                        │               │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ batch_normalization_5           │ (None, 14, 14, 64)     │           256 │
│ (BatchNormalization)            │                        │               │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ leaky_re_lu_5 (LeakyReLU)       │ (None, 14, 14, 64)     │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ conv2d_transpose_5              │ (None, 28, 28, 1)      │         1,600 │
│ (Conv2DTranspose)               │                        │               │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 2,330,944 (8.89 MB)
 Trainable params: 2,305,472 (8.79 MB)
 Non-trainable params: 25,472 (99.50 KB)

Создаем “случайный шум”, с которого всё начинается для Генератора.

Важный момент: Генератор не формирует шум. Шум создается разработчиком/системой и подается на вход генератору.

noise = tf.random.normal([1, 100])
noise

<tf.Tensor: shape=(1, 100), dtype=float32, numpy=
array([[-1.2752672 , -0.31896377, -1.621226  , -0.07633732, -1.1005756 ,
        -0.8244959 ,  0.32265383, -1.3580662 ,  0.81300926,  1.3841189 ,
         1.1405385 ,  1.3428733 , -0.20784518,  0.2218569 , -0.80084634,
        -0.51266044, -0.5123262 ,  1.2493849 , -0.41784754,  0.11716219,
         0.75289106, -0.04998856, -0.10687224,  0.15446882,  0.23294541,
        -0.45333463,  0.29856005, -2.002146  ,  1.035649  ,  0.00998143,
        -1.4422241 , -0.4550751 ,  0.24101041,  0.3818386 ,  0.8918707 ,
         0.3421659 , -1.0747958 ,  0.07026866,  0.92490923,  0.05733351,
        -1.83129   ,  0.07838591, -1.9661248 ,  0.67199177,  0.52293086,
        -0.7199154 ,  1.1893344 ,  1.3752289 , -0.6383991 ,  0.00620717,
         2.937654  , -0.08155467, -0.04186288,  0.2946215 ,  0.08486137,
         0.40340146,  0.31229848,  0.8557363 , -1.2216685 , -0.8172701 ,
        -0.5734942 ,  1.3174502 ,  1.0580747 , -2.3497086 ,  0.331117  ,
        -0.92475957,  2.0144198 , -0.26455778,  1.0036913 ,  0.08436549,
         0.9257641 , -0.35675535, -0.8078421 , -0.06051978,  2.069581  ,
        -0.24463454, -0.8282004 ,  1.6013093 ,  1.0182651 ,  0.87454027,
        -0.27339602,  1.0042901 ,  0.21967338, -2.185581  , -0.562119  ,
         0.5870542 , -0.5030319 ,  0.8525667 , -0.30234253,  1.629835  ,
         1.3273429 , -0.3319834 , -0.37302145,  0.6386749 , -1.4167624 ,
         1.1601201 , -0.89931196,  0.10231236, -0.5188213 ,  0.645322  ]],
      dtype=float32)>

Открываем “шумную” картинку

generated_image = generator(noise, training = False)
generated_image.shape

Подготовка генератора завершена!

Дискриминатор – пытается угадать, какая картинка из настоящего набора данных, а какую только что нарисовали.

Строим сеть дискриминатора :

def build_discriminator():
  network = tf.keras.Sequential()

  # 14x14x64
  network.add(layers.Conv2D(64, (5,5), strides=(2,2), padding='same', input_shape=[28,28,1]))
  network.add(layers.LeakyReLU())
  network.add(layers.Dropout(0.3))

  # 7x7x128
  network.add(layers.Conv2D(128, (5,5), strides=(2,2), padding='same'))
  network.add(layers.LeakyReLU())
  network.add(layers.Dropout(0.3))

  network.add(layers.Flatten())
  network.add(layers.Dense(1))

  network.summary()

  return network

Смотрим сеть дискриминатора :

discriminator = build_discriminator()

-
Model: "sequential_1"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 conv2d (Conv2D)             (None, 14, 14, 64)        1664      
                                                                 
 leaky_re_lu_3 (LeakyReLU)   (None, 14, 14, 64)        0         
                                                                 
 dropout (Dropout)           (None, 14, 14, 64)        0         
                                                                 
 conv2d_1 (Conv2D)           (None, 7, 7, 128)         204928    
                                                                 
 leaky_re_lu_4 (LeakyReLU)   (None, 7, 7, 128)         0         
                                                                 
 dropout_1 (Dropout)         (None, 7, 7, 128)         0         
                                                                 
 flatten (Flatten)           (None, 6272)              0         
                                                                 
 dense_1 (Dense)             (None, 1)                 6273      
                                                                 
=================================================================
Total params: 212,865
Trainable params: 212,865
Non-trainable params: 0
_________________________________________________________________

Подготовка дискриминатора завершена!

Let’s train

X_train
epochs = 100
noise_dim = 100
num_images_to_generate = 16

@tf.function
def train_steps(images):
  noise = tf.random.normal([batch_size, noise_dim])
  with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
    generated_images = generator(noise, training = True)

    expected_output = discriminator(images, training = True)
    fake_output = discriminator(generated_images, training = True)

    gen_loss = generator_loss(fake_output)
    disc_loss = discriminator_loss(expected_output, fake_output)

  gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
  gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

  generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
  discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

test_images = tf.random.normal([num_images_to_generate, noise_dim])

test_images.shape

60000 / 256

def train(dataset, epochs, test_images):
  for epoch in range(epochs):
    for image_batch in dataset:
      #print(image_batch.shape)
      train_steps(image_batch)

    print('Epoch: ', epoch + 1)
    generated_images = generator(test_images, training = False)
    fig = plt.figure(figsize=(10,10))
    for i in range(generated_images.shape[0]):
      plt.subplot(4,4,i+1)
      plt.imshow(generated_images[i, :, :, 0] * 127.5 + 127.5, cmap='gray')
      plt.axis('off')
    plt.show()

train(X_train, epochs, test_images)

2 нейросети соревнуются друг с другом, и благодаря этой борьбе “художник” (Генератор) становится просто гениальным подражателем!

Ниже представлены шаги эпох при обучении нейросети

2. Knowledge Distillation — Дистилляция знаний

Чтобы дистиллировать знания из одной модели в другую, мы берем предобученную teacher-модель, обученную на определенной задаче (в данном случае — классификация изображений), и инициализируем student-модель со случайными весами для обучения на той же задаче классификации изображений.

Перейдем к примеру :

Используем модель merve/beans-vit-224 в качестве teacher-модели. Это модель классификации изображений, основанная на google/vit-base-patch16-224-in21k, дообученная на наборе данных beans (бобовые). Мы будем дистиллировать эту модель в случайно инициализированную MobileNetV2.

Загружаем датасет :

from datasets import load_dataset

dataset = load_dataset("beans")

Подготавливаем данные для работы с teacher-моделью :

from transformers import AutoImageProcessor
teacher_processor = AutoImageProcessor.from_pretrained("merve/beans-vit-224")

def process(examples):
    processed_inputs = teacher_processor(examples["image"])
    return processed_inputs

processed_datasets = dataset.map(process, batched=True)

По сути, мы хотим, чтобы student-модель (случайно инициализированный MobileNet) имитировала teacher-модель (дообученный Vision Transformer). Для достижения этой цели мы сначала получаем данные на выходе как от teacher-модели, так и от student-модели.

from transformers import TrainingArguments, Trainer, infer_device
import torch
import torch.nn as nn
import torch.nn.functional as F

class ImageDistilTrainer(Trainer):
    def __init__(self, teacher_model=None, student_model=None, temperature=None, lambda_param=None,  *args, **kwargs):
        super().__init__(model=student_model, *args, **kwargs)
        self.teacher = teacher_model
        self.student = student_model
        self.loss_function = nn.KLDivLoss(reduction="batchmean")
        device = infer_device()
        self.teacher.to(device)
        self.teacher.eval()
        self.temperature = temperature
        self.lambda_param = lambda_param

    def compute_loss(self, student, inputs, return_outputs=False):
        student_output = self.student(**inputs)

        with torch.no_grad():
          teacher_output = self.teacher(**inputs)


        soft_teacher = F.softmax(teacher_output.logits / self.temperature, dim=-1)
        soft_student = F.log_softmax(student_output.logits / self.temperature, dim=-1)


        distillation_loss = self.loss_function(soft_student, soft_teacher) * (self.temperature ** 2)


        student_target_loss = student_output.loss

      
        loss = (1. - self.lambda_param) * student_target_loss + self.lambda_param * distillation_loss
        return (loss, student_output) if return_outputs else loss

Логинимся в Hugging Face (тут понадобится ваш токен доступа из HF). Через Trainer мы можем выгрузить нашу модель в Hugging Face Hub

from huggingface_hub import notebook_login

notebook_login()

Теперь установим параметры обучения (TrainingArguments), модель-учитель и модель-ученик

from transformers import AutoModelForImageClassification, MobileNetV2Config, MobileNetV2ForImageClassification

repo_name = "ИМЯ ВАШЕГО РЕПОЗИТОРИЯ"

training_args = TrainingArguments(
    output_dir="my-awesome-model",
    num_train_epochs=30,
    fp16=True,
    logging_dir=f"{repo_name}/logs",
    logging_strategy="epoch",
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    report_to="tensorboard",
    push_to_hub=True,
    hub_strategy="every_save",
    hub_model_id=repo_name,
    )

num_labels = len(processed_datasets["train"].features["labels"].names)

# initialize models
teacher_model = AutoModelForImageClassification.from_pretrained(
    "merve/beans-vit-224",
    num_labels=num_labels,
    ignore_mismatched_sizes=True
)

# training MobileNetV2 from scratch
student_config = MobileNetV2Config()
student_config.num_labels = num_labels
student_model = MobileNetV2ForImageClassification(student_config)

После обучения модель будет доступна по ссылке huggingface.co/твой-username/my-model 

Мы можем применить функцию compute_metrics, чтобы оценить нашу модель на тестовой выборке. Данная функция будет задействована во время тренировочного процесса для расчета точности и F1-score модели.

import evaluate
import numpy as np

accuracy = evaluate.load("accuracy")

"""compute_metrics — это кастомная функция,
которая автоматически вызывается Trainer'ом 
во время обучения для мониторинга качества модели."""

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    acc = accuracy.compute(references=labels, predictions=np.argmax(predictions, axis=1))
    return {"accuracy": acc["accuracy"]}

Перейдем к созданию экземпляра Trainer с определенными нами параметрами обучения

from transformers import Trainer

class CompatibleImageDistilTrainer(ImageDistilTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        # Ignore num_items_in_batch and other unexpected kwargs
        return super().compute_loss(model, inputs, return_outputs)

# Use the compatible trainer instead
trainer = CompatibleImageDistilTrainer(
    student_model=student_model,
    teacher_model=teacher_model,
    args=training_args,
    train_dataset=processed_datasets["train"],
    eval_dataset=processed_datasets["validation"],
    data_collator=data_collator,
    processing_class=teacher_processor,
    compute_metrics=compute_metrics,
    temperature=5,
    lambda_param=0.5
)

Следующий шаг после этой настройки — запуск процесса дистилляции, где student будет учиться у teacher

Let’s train)

trainer.train()

Запустим тест и посмотрим на результат

trainer.evaluate(processed_datasets["test"])

Точность: ~ 68%, это на ~5% лучше, чем MobileNet обученный с нуля (~63%).

{'eval_loss': 0.6835939288139343,
 'eval_accuracy': 0.6875,
 'eval_runtime': 13.9354,
 'eval_samples_per_second': 9.185,
 'eval_steps_per_second': 1.148,
 'epoch': 30.0}

Дистилляция сработала успешно! Student model действительно научилась у teacher model. Модель стала значительно лучше благодаря передаче знаний от teacher model к student model.

Дистилляция знаний — это AI-движок для передачи мудрости. Она берёт знания большой, медленной, но умной модели-учителя и упаковывает их в маленькую, быструю модель-ученика — как будто опытный мастер передаёт свои секреты подмастерью.

Ссылка на Google Colab – https://colab.research.google.com/drive/1s4IrEe6JyXpOhpny7UvKwnRPQypr8RRs

3. Reinforcement learning – обучение с подкреплением

Обучение с подкреплением — это про автономность и адаптацию и в условиях неопределенности. RL лежит в основе беспилотных автомобилей, игровых ИИ и роботов — создавая не просто «умные алгоритмы», а самостоятельных рисерчеров, способных к настоящей стратегии и росту через серию проб и ошибок.

Возьмем к примеру обучение с подкреплением на Python с использованием библиотекиgymnasium (современная версия OpenAI Gym).

Ставим необходимые библиотеки и импортируем их

pip install gymnasium numpy

import gymnasium as gym
import numpy as np
import random

Создаем среду FrozenLake – упрощенная игра “найди выход из замерзшего озера”

env = gym.make('FrozenLake-v1', desc=None, map_name="4x4", is_slippery=False)

Инициализируем таблицу (состояния × действия)

# 16 состояний (4x4 клетки) × 4 действия (влево, вниз, вправо, вверх)
q_table = np.zeros([env.observation_space.n, env.action_space.n])

# Параметры обучения
learning_rate = 0.1
discount_factor = 0.99  # Насколько ценим будущие награды
epsilon = 1.0  # Вероятность случайного действия (exploration)
epsilon_decay = 0.999
min_epsilon = 0.01
episodes = 1000

Среда (Environment)

FrozenLake-v1: Упрощенная игра где агент должен дойти от старта (S) до цели (G), избегая провалов в воду (H)

S – старт, F – замерзшая поверхность, H – провал, G – цель

Начинаем обучение :

for episode in range(episodes):
    state, info = env.reset()
    done = False
    total_reward = 0
    
    while not done:
        # Epsilon-greedy стратегия: с вероятностью epsilon выбираем случайное действие
        if random.uniform(0, 1) < epsilon:
            action = env.action_space.sample()  # Случайное действие (exploration)
        else:
            action = np.argmax(q_table[state])  # Лучшее действие (exploitation)
        
        # Выполняем действие в среде
        next_state, reward, terminated, truncated, info = env.step(action)
        done = terminated or truncated
        
        old_value = q_table[state, action]
        next_max = np.max(q_table[next_state])
        
        new_value = old_value + learning_rate * (
            reward + discount_factor * next_max - old_value
        )
        q_table[state, action] = new_value
        
        state = next_state
        total_reward += reward
    
    # Уменьшаем epsilon после каждого эпизода
    epsilon = max(min_epsilon, epsilon * epsilon_decay)
    
    if (episode + 1) % 100 == 0:
        print(f"Эпизод {episode + 1}, Epsilon: {epsilon:.3f}, Награда: {total_reward}")

print("nОбучение завершено!")
print("nИтоговая Q-таблица:")
print(q_table)

"""Получаем вывод

Начинаем обучение...
Эпизод 100, Epsilon: 0.905, Награда: 0.0
Эпизод 200, Epsilon: 0.819, Награда: 0.0
Эпизод 300, Epsilon: 0.741, Награда: 0.0
Эпизод 400, Epsilon: 0.670, Награда: 0.0
Эпизод 500, Epsilon: 0.606, Награда: 0.0
Эпизод 600, Epsilon: 0.549, Награда: 1.0
Эпизод 700, Epsilon: 0.496, Награда: 1.0
Эпизод 800, Epsilon: 0.449, Награда: 1.0
Эпизод 900, Epsilon: 0.406, Награда: 1.0
Эпизод 1000, Epsilon: 0.368, Награда: 1.0

Обучение завершено!

Итоговая Q-таблица:
[[0.94147845 0.95099005 0.93191546 0.94146579]
 [0.94144155 0.         0.76517588 0.79404897]
 [0.25227238 0.92225733 0.18540704 0.37700074]
 [0.44469117 0.         0.09322424 0.02430739]
 [0.9509832  0.96059601 0.         0.94146724]
 [0.         0.         0.         0.        ]
 [0.         0.97919907 0.         0.56246919]
 [0.         0.         0.         0.        ]
 [0.96045605 0.         0.970299   0.95097178]
 [0.95927825 0.97900787 0.9801     0.        ]
 [0.96991831 0.99       0.         0.95961833]
 [0.         0.         0.         0.        ]
 [0.         0.         0.         0.        ]
 [0.         0.89988563 0.98996049 0.86611194]
 [0.97278094 0.98831215 1.         0.97813866]
 [0.         0.         0.         0.        ]]

"""

Тестируем обученного агента :

print("nТестируем обученного агента:")
state, info = env.reset()
done = False
steps = 0

while not done and steps < 20:
    action = np.argmax(q_table[state])  # Всегда выбираем лучшее действие
    state, reward, terminated, truncated, info = env.step(action)
    done = terminated or truncated
    steps += 1
    
    env.render()  # Показываем визуализацию
    print(f"Шаг {steps}: Действие {action}, Награда {reward}")
    
    if done:
        if reward > 0:
            print("🎉 Успех! Агент достиг цели!")
        else:
            print("💥 Провал! Агент упал в воду!")

env.close()

"""

Вывод :

Шаг 1: Действие 1, Награда 0.0
 Шаг 2: Действие 1, Награда 0.0
 Шаг 3: Действие 2, Награда 0.0
 Шаг 4: Действие 2, Награда 0.0
 Шаг 5: Действие 1, Награда 0.0
 Шаг 6: Действие 2, Награда 1.0
 🎉 Успех! Агент достиг цели!

"""

После достаточного количества эпизодов агент научится:

• Находить безопасный путь от старта к цели

• Избегать провалов в воду

• Действовать оптимально в каждом состоянии

Этот пример демонстрирует основные принципы RL: агент взаимодействует со средой, получает награды и учится через метод проб и ошибок!

Дополнительная визуализация:

def print_policy(q_table):
    actions = ['←', '↓', '→', '↑']
    policy = []
    
    for state in range(16):
        best_action = np.argmax(q_table[state])
        policy.append(actions[best_action])
    

    for i in range(0, 16, 4):
        print(' '.join(policy[i:i+4]))

print_policy(q_table)

Вывод:
  
  """
↓ ← ↓ ←
↓ ← ↓ ←
→ → ↓ ←
← → → ←
"""

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

Ссылка на Google Colab – https://colab.research.google.com/drive/12Qu0vF6ETfO7s5PiD0u_fePJZ72f8TM5#scrollTo=a2crPdCoDKBy

на этом все) до новых встреч)