Логистическая регрессия для модных футболок MNIST и рубашек

Я написал Логистическая регрессия для Fashion MNIST классифицировать футболку против рубашки.

вот класс

import numpy as np
import time

class LogisticRegression:
    def __init__(self, learning_rate=.05, n_epoch=10, model_w=[], model_b=.0):
        self.learning_rate = learning_rate
        self.n_epoch = n_epoch
        self.model_w = model_w
        self.model_b = model_b
        
    def initialize_params(self, n_x):
        if len(self.model_w) == 0:
            self.model_w = np.random.random((n_x, 1))
    
    # forward_propagation
    # predict_prob
    def activation(self, x):
        z = np.dot(x, self.model_w) + self.model_b
        a = 1 / (1 + np.exp(-z))
        return a
        
    def predict(self, x):
        a = self.activation(x) >= 0.5
        return np.squeeze(1*a)
    
    def evaluate(self, x, y):
        acc = np.count_nonzero(np.squeeze(self.predict(x)) == np.squeeze(y))/len(y)
        return acc
    
    def backward_propagation(self, x, a, y):
        m = len(x)
        dz = a - y
        dw = np.dot(x.T, dz)/m
        db = np.mean(dz)
        return dw, db

    def update_weights(self, dw, db):
        self.model_w -= self.learning_rate * dw
        self.model_b -= self.learning_rate * db
        
    def fit(self, x, y, batch_size=10, shuffle=True, verbose=True):
        start_time = time.time()
        n_x = x.shape[-1]
        self.initialize_params(n_x)
        indices = np.arange(len(x))
        for i in range(self.n_epoch):
            if shuffle:
                np.random.shuffle(indices)
            n_batches = int(len(x)/batch_size)
            batches = np.split(indices[:batch_size*n_batches], n_batches)
            for batch in batches:
                a = self.activation(x[batch])
                dw, db = self.backward_propagation(x[batch], a, y[batch])
                self.update_weights(dw, db)
        acc = self.evaluate(x, y)
        if verbose:
            print('model trained {:.5f} s'.format(time.time() - start_time))
            print('train accuracy: {:7f}'.format(acc))

Вот данные

import tensorflow as tf
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# generate the indices
idx_digit_train = np.argwhere((y_train == 0) | (y_train == 6)).flatten()
idx_digit_test = np.argwhere((y_test == 0) | (y_test == 6)).flatten()

# construct the training set
y_train_mnist = y_train[idx_digit_train].reshape(-1,1)
x_train_mnist = x_train[idx_digit_train]

# construct the test set
y_test_mnist = y_test[idx_digit_test].reshape(-1,1)
x_test_mnist = x_test[idx_digit_test]

# normalization
x_train_mnist = x_train_mnist/255.
x_test_mnist = x_test_mnist/255.

# flatten
x_train_mnist = x_train_mnist.reshape(len(x_train_mnist), -1)
x_test_mnist = x_test_mnist.reshape(len(x_test_mnist), -1)

y_train_mnist[y_train_mnist==6]=1
y_test_mnist[y_test_mnist==6]=1

Обучил модель на данных за одну эпоху

classifier = LogisticRegression(.1, 1)
classifier.fit(x_train_mnist, y_train_mnist)
classifier.evaluate(x_test_mnist, y_test_mnist)

Вот результат

model trained 0.09662 s
train accuracy: 0.828583
0.8125

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