MNISTの0〰9数字を分類するコードで学習が進みません。

Question

前提

初心者です。
MNISTの0〰9数字を分類するコードをスクラッチで書いています。

実現したいこと

softmax関数の計算がおかしいようで学習が進ません。
class SoftmaxCrossEntropyLoss()の中で使っています。
softmaxの書き方が悪いのか、層の計算が間違っているのかも分かりません。
何んとなく怪しい箇所は、#怪しい と記載しました。

発生している問題・エラーメッセージ

RuntimeWarning: invalid value encountered in subtract
  x = x-np.max(x,axis=0)

該当のソースコード

import csv
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import cv2
from sklearn import datasets
from sklearn.model_selection import train_test_split
%matplotlib inline
np.random.seed(seed=0)


#怪しい
def softmax(x):
    x = x.T
    x = x-np.max(x,axis=0)
    x = np.exp(x)/np.sum(np.exp(x))
    return x.T

# mnistデータセット
mnist = datasets.fetch_openml('mnist_784', as_frame=False)
# 画像とラベルを取得
X, T = mnist.data, mnist.target
# 訓練データとテストデータに分割
X_train, X_test, T_train, T_test = train_test_split(X, T, test_size=0.2)

T_train = np.eye(10)[T_train.astype("int")]
T_test = np.eye(10)[T_test.astype("int")]

def cross_entropy_error(t, y):
    delta = 1e-8
    error = -np.mean(t * np.log(y + delta))
    return error

class SoftmaxCrossEntropyLoss():
    def __init__(self):
        self.y = None
        self.t = None
        self.loss = None
        
    def __call__(self, t, y):
        self.y = softmax(y)
        self.t = t.copy()
        self.loss = cross_entropy_error(self.t, self.y)
        return self.loss
    
    def backward(self):
        batch_size = self.t.shape[0]
        dy = self.y - self.t
        dy /= batch_size
        return dy

class FullyConnectedLayer():
    def __init__(self, input_shape, output_shape):
        self.w = np.random.randn(input_shape, output_shape) * 0.01
        self.b = np.zeros(output_shape, dtype=np.float)
        self.x = None
        self.dw = None
        self.db = None

#怪しい        
    def __call__(self, x):
        self.x = x
        out = np.dot(x,self.w)+self.b
        return out
#怪しい    
    def backward(self, dout):
        dx = np.dot(dout,np.transpose(self.w))
        batch_size = dx.shape[0]
        self.dw = np.dot(np.transpose(self.x),dout)
        self.db =  np.sum(dout,axis=0)
        return dx

class ReLU():
    def __init__(self):
        self.mask = None
#怪しい
    def __call__(self, x):
        self.mask = (x <= 0)
        out = x.copy() 
        out[self.mask]=0 
        return out
#怪しい
    def backward(self, dout):
        dout[self.mask]=0 
        dx = dout

        return dx

class MLP_classifier():

    def __init__(self):
        '''
        
        x -> fc(784, 256) -> relu -> fc(256, 256) -> relu -> fc(256, 10) -> out
        '''
        
        # 層
        self.fc1 = FullyConnectedLayer(784, 256)
        self.relu1 = ReLU()
        self.fc2 = FullyConnectedLayer(256, 256)
        self.relu2 = ReLU()
        self.fc3 = FullyConnectedLayer(256, 10)
        self.out = None
        
        # 損失関数の定義
        self.criterion = SoftmaxCrossEntropyLoss()

    def forward(self, x):
        '''
        順伝播
        '''
        
        x = self.relu1(self.fc1(x))
        x = self.relu2(self.fc2(x))
        self.out = self.fc3(x)
        

        return self.out

    def backward(self, t):
        '''
        逆伝播
        '''
        
        # 誤差を計算
        loss = self.criterion(t, self.out)
        # 勾配を逆伝播
        d = self.criterion.backward()
        d = self.fc3.backward(d)
        d = self.relu2.backward(d)
        d = self.fc2.backward(d)
        d = self.relu1.backward(d)
        d = self.fc1.backward(d)
        
        return loss

    def optimize_GradientDecent(self, lr):
        '''
        勾配降下法による全層のパラメータの更新
        '''
        for fc in [self.fc1, self.fc2, self.fc3]:
            fc.w -= lr * fc.dw
            fc.b -= lr * fc.db

# モデルの宣言
model = MLP_classifier()

# 学習率
lr = 0.005
# 学習エポック数
n_epoch = 20


for n in range(n_epoch):
    # 訓練

    y = model.forward(X_train)
    loss = model.backward(T_train)
    model.optimize_GradientDecent(lr)
    
    # テスト
    y = model.forward(X_test)
    test_loss = model.backward(T_test)
    pred = softmax(y)
    accuracy = np.mean(np.equal(np.argmax(y, axis=1), np.argmax(T_test, axis=1)))
    print(f'EPOCH {n + 1} | TRAIN LOSS {loss:.5f} | TEST LOSS {test_loss:.5f} | ACCURACY {accuracy:.2%}')
classification_accuracy = accuracy

試したこと

エラーがdef softmax()の中なので、その中はいくつか試しました。

補足情報（FW/ツールのバージョンなど）

初心者ですので、質問の仕方が不適切かと思いますが、
助けてください。

Comments

[meg_]

まずは順伝播の計算が合っていることを単純なデータで確認するのが先かと思います。あと、エラーではなくウォーニングかと思いますが。

[1T2R3M4]

https://qiita.com/kobaq/questions/f2b8aa13f6d4d54669e9
以下ご対応ください。
https://teratail.com/help#posted-otherservice

[jbpb0]

https://teratail.com/help#posted-otherservice
に、
「解決した際には必ずteratail及びすべての投稿に解決した旨と、どのように解決したかを記載してください。」
と書かれてますので、
https://qiita.com/kobaq/questions/f2b8aa13f6d4d54669e9
の回答で解決したなら、その内容をこちらの回答に書いて「自己解決」にしてください

[kobahot]

ご指導ありがとうございます。ルールを把握しておらず申し訳ありません。
正直、解決したか自信ありませんが、学習は完了できたので、その内容を投稿しておきます。

Answer 1

https://qiita.com/segavvy/items/bdad9fcda2f0da918e7c
↑を参考に、keepdims=Trueを追加し、np.maxで次元が減らないように修正した所、学習が完了しました。