Реализация нейронной сети “с нуля”. Добавление слоев сломало код каким-то образом.

Я читаю книгу под названием “Глубокое обучение на Python”. Я нахожусь на этапе, где автор объясняет реализацию нейронной сети в “сыром” TensorFlow. То есть, используя минимальное количество Keras.

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

Посмотрите на результаты:

Эпоха 0
потеря на пакете 0: 12.23
потеря на пакете 100: 14.48
потеря на пакете 200: 14.48
потеря на пакете 300: 13.10
потеря на пакете 400: 14.98
Эпоха 1
потеря на пакете 0: 13.60
потеря на пакете 100: 14.48
потеря на пакете 200: 14.48
потеря на пакете 300: 13.10
потеря на пакете 400: 14.98
Эпоха 2
потеря на пакете 0: 13.60
потеря на пакете 100: 14.48
потеря на пакете 200: 14.48
потеря на пакете 300: 13.10
потеря на пакете 400: 14.98
Эпоха 3
потеря на пакете 0: 13.60
потеря на пакете 100: 14.48
потеря на пакете 200: 14.48
потеря на пакете 300: 13.10
потеря на пакете 400: 14.98
Эпоха 4
потеря на пакете 0: 13.60
потеря на пакете 100: 14.48
потеря на пакете 200: 14.48
потеря на пакете 300: 13.10
потеря на пакете 400: 14.98
точность: 0.11

По какой-то причине потеря вообще не уменьшается. Мы попадаем в некий цикл. Найдено ли в коде локальное минимум или что-то еще? Вот код:

import tensorflow as tf
import math
import numpy as np

class NaiveDense:
  def __init__(self, input_size, output_size, activation):
    self.activation = activation
    w_shape = (input_size, output_size)
    w_intial_value = tf.random.uniform(w_shape, minval = 0, maxval=1e-1)
    self.W = tf.Variable(w_intial_value)

    b_shape = (output_size,)
    b_initial_value = tf.zeros(b_shape)
    self.b = tf.Variable(b_initial_value)

  # прямой проход
  def __call__(self, inputs):
    return self.activation(tf.matmul(inputs, self.W) + self.b)

  @property
  def weights(self):
    return [self.W, self.b]

class NaiveSequential:
  def __init__(self, layers):
    self.layers = layers

  def __call__(self, inputs):
    x = inputs
    for layer in self.layers:
      x = layer(x)
    return x

  @property
  def weights(self):
    weights = []
    for layer in self.layers:
      weights += layer.weights
    return weights

model = NaiveSequential([
    NaiveDense(input_size = 28 * 28, output_size=512, activation=tf.nn.relu),
    NaiveDense(input_size = 512, output_size=256, activation=tf.nn.softplus),
    NaiveDense(input_size = 256, output_size=10, activation=tf.nn.softmax)
  ])

class BatchGenerator:
  def __init__(self, images, labels, batch_size = 128):
    assert len(images) == len(labels)
    self.index = 0
    self.images = images
    self.labels = labels
    self.batch_size = batch_size
    self.num_batches = math.ceil(len(images) / batch_size)

  def next(self):
    images = self.images[self.index : self.index + self.batch_size]
    labels = self.labels[self.index : self.index + self.batch_size]
    self.index += self.batch_size
    return images, labels

def one_training_step(model, images_batch, labels_batch):
  with tf.GradientTape() as tape:
    predicitions = model(images_batch)
    per_sample_losses = tf.keras.losses.sparse_categorical_crossentropy(labels_batch, predicitions)
    average_loss = tf.reduce_mean(per_sample_losses)
  gradients = tape.gradient(average_loss, model.weights)
  update_weights(gradients, model.weights)
  return average_loss

learning_rate = 1e-3

def update_weights(gradients, weights):
  for g, w in zip(gradients, weights):
    w.assign_sub(g * learning_rate)

def fit(model, images, labels, epochs, batch_size=128):
  for epoch_counter in range(epochs):
    print(f"Эпоха {epoch_counter}")
    batch_generator = BatchGenerator(images, labels)
    for batch_counter in range(batch_generator.num_batches):
      images_batch, labels_batch = batch_generator.next()
      loss = one_training_step(model, images_batch, labels_batch)
      if batch_counter % 100 == 0:
        print(f"потеря на пакете {batch_counter}: {loss:.2f}")

from tensorflow.keras.datasets import mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

train_images = train_images.reshape((60000, 28 * 28))
train_images = train_images.astype("float32") / 255
test_images = test_images.reshape((10000, 28 * 28))
test_images = test_images.astype("float32") / 255

fit(model, train_images, train_labels, epochs=5, batch_size = 128)

predictions = model(test_images)
precictions = predictions.numpy()
predicted_labels = np.argmax(predictions, axis = 1)
matches = predicted_labels == test_labels
print(f"точность: {matches.mean():.2f}")

Я пытался добавить слой в Keras, и я получил желаемый результат. То есть точность увеличивается с 0.9793 до 0.9811.

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.datasets import mnist

(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

model = keras.Sequential([
     layers.Dense(512, activation = "relu"),
     layers.Dense(256, activation = "relu"),
     layers.Dense(10, activation = "softmax")     
])

model.compile(optimizer="rmsprop",
              loss = "sparse_categorical_crossentropy",
              metrics = ["accuracy"])

train_images = train_images.reshape((60000, 28 * 28))
train_images = train_images.astype("float32") / 255
test_images = test_images.reshape((10000, 28 * 28))
test_images = test_images.astype("float32") / 255

model.fit(train_images, train_labels, epochs=5, batch_size=128)

test_loss, test_acc = model.evaluate(test_images, test_labels)
print(f"Тестовая точность: {test_acc}")

Я, возможно, упустил что-то в расчете градиента в коде TensorFlow? Я знаю, что активации внутреннего слоя различаются в каждом случае, но я пробовал использовать одну и ту же активацию, и в результате получаю этот “цикл” в градиенте.