Tensorflow 2 入門 AND回路を学習させる

Tensorflowは以下のサイトにチュートリアルがありますが、難しく感じたのでより単純な学習をさせてみました。
www.tensorflow.org

・環境
　python 3.7.7
　Tensor Flow 2.2
　Anaconda 3

・やること
　AND回路入力の1と０、そして出力結果を学習させる。

AND回路

AND回路 真理値表

・学習用プログラム

# TensorFlow と tf.keras のインポート
import tensorflow as tf
from tensorflow import keras

# ヘルパーライブラリのインポート
import numpy as np

################################################
#学習用とテスト用データ
################################################
train_x=[[0,0],[0,1],[1,0],[1,1]]   #トレーニング用入力
train_y=[[0],[0],[0],[1]]           #トレーニング用出力

test_x=[[0,0],[0,1],[1,0],[1,1]]    #テスト用入力
test_y=[[0],[0],[0],[1]]            #テスト用出力

################################################
#モデル構築
################################################
#データフォーマット変換
model = keras.Sequential([
    keras.layers.Flatten(input_shape=(1, 2)),
    keras.layers.Dense(128, activation='relu'),   #128個のニューロン
    keras.layers.Dense(2, activation='softmax') #出力数が2個なので2と指定（出力数の設定）
])

model.summary() # モデルの構造を表示

model.compile(optimizer='adam', 
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(train_x, train_y, epochs=1000)         #トレーニング開始

################################################
#正答率の評価
################################################
test_loss, test_acc = model.evaluate(test_x,  test_y, verbose=2)
print('\nTest accuracy:', test_acc)

################################################
#予測
################################################
predictions = model.predict(test_x)
print(predictions[0])                   #各確率を表示
print(np.argmax(predictions[0]))        #一番確率が高い物のラベル表示

################################################
# モデル全体を SavedModel として保存
################################################
model.save('my_model.h5')

学習させるものとしては4パターンしかないためトレーニングデータとテストデータは全く同じにしています。

train_x=[[0,0],[0,1],[1,0],[1,1]] がAND回路の入力データとなり、
train_y=[[0],[0],[0],[1]]が出力データとなります

model = keras.Sequential([
keras.layers.Flatten(input_shape=(1, 2)),
keras.layers.Dense(128, activation='relu'),
keras.layers.Dense(2, activation='softmax')
])
のinput_shape=(1, 2)は入力のデータ配列を表しており、AND回路の入力データは1列２行なので(1,2)にしています。
入力が画像データだとinput_shape=(縦ピクセル数, 横ピクセル数)となります。
keras.layers.Dense(2, activation='softmax')の2の数字は出力数です。AND回路の出力結果は0か1の結果しかないので2です。
画像の学習で結果が、猫,犬,パンダと３つある場合は3をいれます。そうすると出力結果としては0～2で出力されます。

・学習させたモデルを読み込み学習されているかを確認するプログラム

# TensorFlow と tf.keras のインポート
import tensorflow as tf
from tensorflow import keras

# ヘルパーライブラリのインポート
import numpy as np

test_data=[[0,0],[0,1],[1,0],[1,1]]

#同じモデルを読み込んで、重みやオプティマイザーを含むモデル全体を再作成
new_model = tf.keras.models.load_model('my_model.h5')
# モデルのアーキテクチャを表示
new_model.summary()

predictions = new_model.predict(test_data)

print('結果')
print('0,0=',np.argmax(predictions[0]))
print('0,1=',np.argmax(predictions[1]))
print('1,0=',np.argmax(predictions[2]))
print('1,1=',np.argmax(predictions[3]))

このプログラムは学習用プログラムで作成したモデルmy_model.h5を読み込みちゃんと学習されているかを確認するためのものです。
テストデータは学習させたときと同じものを入れ、
new_model = tf.keras.models.load_model('my_model.h5')でモデルを読み込み
predictions = new_model.predict(test_data)でテストデータを判定し結果をpredictions に入れています。
そして最後に結果が出力されます。