TensorFlow で MNIST その４【TensorFlow.js でトレーニング】
過去のエントリーでは Python でトレーニングしたモデルを変換して TensorFlow.js で使いましたが、 今回は トレーニングから TensorFlow.js で行います。 そのため、 このエントリーのトレーニング用のコード（モデルを構築するなど） を JavaScript (TensorFlow.js) に移植します。
このページでは、コードを簡潔に表現するため、 生成したテンソルの破棄する記述はあえて省略しています。 ご了承ください。 必要なら、 tf.dispose() や tf.tidy() を適宜使用してください。
データセットを準備
トレーニグするにはMNISTの手描き数字データとその正解ラベルを用意する必要があります。
Python 版の keras では以下の１行でそれが用意できていました。
# Load the data and split it between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
ここでは、このデータセットを一旦ファイルに書き出した上で、TensorFlow.js から使うことにします。 細かい説明は省きますが、以下のコードを使います。
Python 3.9.6 を使用。 pip install tensorflow==2.13.0 pillow==10.0.0 済みとする。
download.py
import os
from tensorflow import keras
from PIL import Image
def export_jpg_file(export_dir, x):
image_size_y = x.shape[1]
image_size_x = x.shape[2]
for i in range( len(x) ):
img = Image.new('L', (image_size_x, image_size_y))
data = x[i]
jpg_filename = str(i).zfill(5) + '.jpg'
for rowIndex in range(data.shape[0]):
for colIndex in range(data.shape[1]):
value = data[rowIndex][colIndex]
img.putpixel((colIndex, rowIndex), int(value))
img.save( os.path.join(export_dir, jpg_filename) )
def export_csv_file(export_dir, y):
num_classes = 10
y_onehot = keras.utils.to_categorical(y, num_classes)
for i in range( y_onehot.shape[0] ):
csv_filename = str(i).zfill(5) + '.csv'
csv_file_path = os.path.join(export_dir, csv_filename)
a = ','.join( list( map(lambda it: str(int(it)), y_onehot[i]) ) )
with open(csv_file_path, "w") as f:
f.write(a)
train_dir = 'data/train'
train_img_dir = os.path.join(train_dir, 'img')
train_label_dir = os.path.join(train_dir, 'label')
test_dir = 'data/test'
test_img_dir = os.path.join(test_dir, 'img')
test_label_dir = os.path.join(test_dir, 'label')
for it in [train_img_dir, train_label_dir, test_img_dir, test_label_dir]:
os.makedirs(it, exist_ok=True)
# Load the data and split it between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
# Export images
export_jpg_file(train_img_dir, x_train[:10])
export_jpg_file(test_img_dir, x_test[:3])
# Export answers
export_csv_file(train_label_dir, y_train[:10])
export_csv_file(test_label_dir, y_test[:3])
python download.py して実行すると ./data/ 以下にデータセット（の一部）が用意されます。
.
└── data
├── test
│ ├── img
│ │ ├── 00000.jpg
│ │ ├── 00001.jpg
│ │ └── 00002.jpg
│ └── label
│ ├── 00000.csv
│ ├── 00001.csv
│ └── 00002.csv
└── train
├── img
│ ├── 00000.jpg
│ ├── 00001.jpg
│ ├── 00002.jpg
│ ├── 00003.jpg
│ ├── 00004.jpg
│ ├── 00005.jpg
│ ├── 00006.jpg
│ ├── 00007.jpg
│ ├── 00008.jpg
│ └── 00009.jpg
└── label
├── 00000.csv
├── 00001.csv
├── 00002.csv
├── 00003.csv
├── 00004.csv
├── 00005.csv
├── 00006.csv
├── 00007.csv
├── 00008.csv
└── 00009.csv
たとえば ./data/train/img/00007.jpg は以下の画像です。
そして、これに対する正解ラベル(one hot encoding) である ./data/train/label/00007.csv は以下の通り
0,0,0,1,0,0,0,0,0,0
0-index で数えて 3番目が 1 になっています。（手描き数字の内容が 3 なので...）
今は内容を見るために 以下のコードの x_train[:10] のようにしてデータセットとして保存する数を制限していました。
# Export images
export_jpg_file(train_img_dir, x_train[:10])
export_jpg_file(test_img_dir, x_test[:3])
# Export answers
export_csv_file(train_label_dir, y_train[:10])
export_csv_file(test_label_dir, y_test[:3])
それでは実際にトレーニングするために、 このコードを以下のように修正して、すべてのデータをローカルに保存します。
# Export images
export_jpg_file(train_img_dir, x_train)
export_jpg_file(test_img_dir, x_test)
# Export answers
export_csv_file(train_label_dir, y_train)
export_csv_file(test_label_dir, y_test)
これで ./data/ 以下に train 画像 60000 件、test 画像 10000 件を用意することができました。
TensorFlow.js
ここからは Python のことは忘れて JavaScript のみを使います。 TensorFlow.js を使うので、Node.js が必要になります。
以下のような環境で作業しています。
$ node -v
v18.12.1
4 $ npm -v
8.19.2
それでは npm で TensorFlow.js をインストールします。
$ npm init -y
$ npm install @tensorflow/tfjs-node@4.10.0
$ touch train.js
それでは train.js にトレーニング用コードを書きます。
データセットをロード
最初は Python keras のコード、以下の部分を JavaScript に移植します。
# Load the data and split it between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train は ./data/train/img/*.jpg が shape [60000, 28, 28, 1] のテンソルになっているものです。
const fs = require('fs')
const path = require('path')
const tf = require('@tensorflow/tfjs-node')
const trainImgDir = 'data/train/img'
const trainImgFilenames = fs.readdirSync(trainImgDir)
const trainImgFiles = trainImgFilenames.map((filename)=> path.join(trainImgDir, filename))
console.log(trainImgFiles)
node train.js として実行すると以下のようになります。
$ node train.js
[
'data/train/img/00000.jpg',
'data/train/img/00001.jpg',
'data/train/img/00002.jpg',
'data/train/img/00003.jpg',
'data/train/img/00004.jpg',
'data/train/img/00005.jpg',
'data/train/img/00006.jpg',
'data/train/img/00007.jpg',
'data/train/img/00008.jpg',
'data/train/img/00009.jpg',
'data/train/img/00010.jpg',
... more items
]
これで、ファイルパスのリストを得たので、あとは 所定のテンソル（ [ 60000, 28, 28, 1 ]）に変換するだけです。
一つの jpg 画像をテンソルに変換する関数 toImgTensor を定義します。
const toImgTensor = (jpgFile)=> {
const image = fs.readFileSync(jpgFile)
const tensor = tf.node.decodeImage(image, 1)
return tensor.div(255) // 0..1 の範囲に値をスケール.
}
これを 60000 件繰り返して tf.stack() します。
const xTrain = tf.stack( trainImgFiles.map( toImgTensor ) )
console.log( `xTrain => [ ${xTrain.shape} ]` )
実行して、意図した shape になっているか確認します。
$ node train.js
xTrain => [ 60000, 28, 28, 1 ]
問題ありません。
次に正解ラベルのCSVファイルから y_train に相当するテンソルを用意します。
const trainLabelDir = 'data/train/label'
const trainLabelFilenames = fs.readdirSync(trainLabelDir)
const trainLabelFiles = trainLabelFilenames.map((filename)=> path.join(trainLabelDir, filename))
ここまでは、画像の場合とほぼ同じです。（対象ディレクトリが data/train/label に変わったにすぎない。）
正解ラベルが記載されたCSVファイルを読んで １階のテンソルに変換する関数 toLabelTensor を定義。
const toLabelTensor = (csvFile)=> {
const text = fs.readFileSync(csvFile, 'utf8')
const values = text.split(/,/).map( (value)=> parseInt(value) )
return tf.tensor(values, null, 'float32')
}
あとはスタックするだけです。
const yTrain = tf.stack( trainLabelFiles.map( toLabelTensor ) )
console.log( `yTrain => [ ${yTrain.shape} ]` )
これで y_train に相当する TensorFlow.js のテンソル yTrain が用意できました。
$ node train.js
xTrain => [ 60000, 28, 28, 1 ]
yTrain => [ 60000, 10 ]
x_test, y_test に相当するテンソルも同様にして作ります。 データをロードする部分の完成したコードは以下の通りです。
train.js
const fs = require('fs')
const path = require('path')
const tf = require('@tensorflow/tfjs-node')
const loadData = ()=>{
const toImgTensor = (jpgFile)=> {
const image = fs.readFileSync(jpgFile)
const tensor = tf.node.decodeImage(image, 1)
return tensor.div(255) // 0..1 の範囲に値をスケール.
}
const toLabelTensor = (csvFile)=> {
const text = fs.readFileSync(csvFile, 'utf8')
const values = text.split(/,/).map( (value)=> parseInt(value) )
return tf.tensor(values, null, 'float32')
}
// 1) train
const trainImgDir = 'data/train/img'
const trainImgFilenames = fs.readdirSync(trainImgDir)
const trainImgFiles = trainImgFilenames.map((filename)=> path.join(trainImgDir, filename))
const xTrain = tf.stack( trainImgFiles.map( toImgTensor ) )
const trainLabelDir = 'data/train/label'
const trainLabelFilenames = fs.readdirSync(trainLabelDir)
const trainLabelFiles = trainLabelFilenames.map((filename)=> path.join(trainLabelDir, filename))
const yTrain = tf.stack( trainLabelFiles.map( toLabelTensor ) )
// 2) test
const testImgDir = 'data/test/img'
const testImgFilenames = fs.readdirSync(testImgDir)
const testImgFiles = testImgFilenames.map((filename)=> path.join(testImgDir, filename))
const xTest = tf.stack( testImgFiles.map( toImgTensor ) )
const testLabelDir = 'data/test/label'
const testLabelFilenames = fs.readdirSync(testLabelDir)
const testLabelFiles = testLabelFilenames.map((filename)=> path.join(testLabelDir, filename))
const yTest = tf.stack( testLabelFiles.map( toLabelTensor ) )
return {
xTrain: xTrain,
yTrain: yTrain,
xTest: xTest,
yTest: yTest}
}
const dataset = loadData()
const xTrain = dataset.xTrain
const yTrain = dataset.yTrain
const xTest = dataset.xTest
const yTest = dataset.yTest
console.log( `xTrain => [ ${xTrain.shape} ]` )
console.log( `yTrain => [ ${yTrain.shape} ]` )
console.log( `xTest => [ ${xTest.shape} ]` )
console.log( `yTest => [ ${yTest.shape} ]` )
実行します。
$ node train.js
xTrain => [ 60000,28,28,1 ]
yTrain => [ 60000,10 ]
xTest => [ 10000,28,28,1 ]
yTest => [ 10000,10 ]
できました。
モデルの構築
モデルは Python keras では以下のように記述していました。
model = keras.Sequential(
[
keras.Input(shape=input_shape),
layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Flatten(),
layers.Dropout(0.5),
layers.Dense(num_classes, activation="softmax"),
]
)
これは model.add() を使って表現すると以下のようになります。
model = Sequential()
model.add(layers.Conv2D(32, kernel_size=(3,3), activation="relu", input_shape=input_shape))
model.add(layers.MaxPooling2D(pool_size=(2,2)))
model.add(layers.Conv2D(64, kernel_size=(3,3), activation="relu"))
model.add(layers.MaxPooling2D(pool_size=(2,2)))
model.add(layers.Flatten())
model.add(layers.Dropout(0.5))
model.add(layers.Dense(num_classes, activation="softmax"))
これを TensorFlow.js に移植します。
const createModel = (numClasses, inputShape)=> {
const model = tf.sequential()
// model.add(layers.Conv2D(32, kernel_size=(3,3), activation="relu", input_shape=input_shape))
model.add(
tf.layers.conv2d({
filters: 32,
kernelSize: 3,
strides: 1,
padding: 'valid',
activation: 'relu',
inputShape: inputShape
})
)
// model.add(layers.MaxPooling2D(pool_size=(2,2)))
model.add(tf.layers.maxPooling2d({ poolSize: [2, 2] }))
// model.add(layers.Conv2D(64, kernel_size=(3,3), activation="relu"))
model.add(
tf.layers.conv2d({
filters: 64,
kernelSize: 3,
strides: 1,
padding: 'valid',
activation: 'relu',
})
)
// model.add(layers.MaxPooling2D(pool_size=(2,2)))
model.add(tf.layers.maxPooling2d({ poolSize: [2, 2] }))
// model.add(layers.Flatten())
model.add(tf.layers.flatten())
// model.add(layers.Dropout(0.5))
model.add(tf.layers.dropout(0.5))
// model.add(layers.Dense(num_classes, activation="softmax"))
model.add(
tf.layers.dense({
units: numClasses,
activation: 'softmax'
})
)
return model
}
TensorFlow.js に Python 版 keras のレイヤー等が用意されている限りは、機械的な記述変更になります。
const numClasses = 10
const inputShape = [28, 28, 1]
const model = createModel(numClasses, inputShape)
model.summary()
実行してモデルのサマリを見ます。
$ node train.js
__________________________________________________________________________________________
Layer (type) Input Shape Output shape Param #
==========================================================================================
conv2d_Conv2D1 (Conv2D) [[null,28,28,1]] [null,26,26,32] 320
__________________________________________________________________________________________
max_pooling2d_MaxPooling2D1 [[null,26,26,32]] [null,13,13,32] 0
__________________________________________________________________________________________
conv2d_Conv2D2 (Conv2D) [[null,13,13,32]] [null,11,11,64] 18496
__________________________________________________________________________________________
max_pooling2d_MaxPooling2D2 [[null,11,11,64]] [null,5,5,64] 0
__________________________________________________________________________________________
flatten_Flatten1 (Flatten) [[null,5,5,64]] [null,1600] 0
__________________________________________________________________________________________
dropout_Dropout1 (Dropout) [[null,1600]] [null,1600] 0
__________________________________________________________________________________________
dense_Dense1 (Dense) [[null,1600]] [null,10] 16010
==========================================================================================
Total params: 34826
Trainable params: 34826
Non-trainable params: 0
Python 版 keras の model.summary() 結果と一致しているか確認しましょう。
トレーニングを実行
最後にこのモデルとデータセットを使ってトレーニングします。
Python の keras のコードでは以下の通りです。
batch_size = 128
epochs = 15
model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)
これを TensorFlow.js 用に書き直します。
model.compile({
loss: 'categoricalCrossentropy',
optimizer: 'adam',
metrics: ['accuracy']
})
await model.fit(trainX, trainY, {
epochs: epochs,
callbacks: printCallback,
batchSize: batchSize,
validationSplit: 0.1
})
これを修正して、トレーニングの評価用に xTest, yTest を使うようにします。 また、トレーニング結果を ./mymodel/ に保存する処理を追加します。
model.compile({
loss: 'categoricalCrossentropy',
optimizer: 'adam',
metrics: ['accuracy']
})
await model.fit(xTrain, yTrain, {
epochs: epochs,
callbacks: printCallback,
batchSize: batchSize,
validationData: [xTest, yTest]
//validationSplit: 0.2
})
const modelDirPath = path.join(__dirname, 'mymodel')
await model.save(`file://${modelDirPath}`)
この部分は、await を使うため async 関数で囲む必要があります。 従って最終的には以下のようになりました。
const train = async (model, dataset, batchSize, epochs, printCallback) => {
const xTrain = dataset.xTrain
const yTrain = dataset.yTrain
const xTest = dataset.xTest
const yTest = dataset.yTest
model.compile({
loss: 'categoricalCrossentropy',
optimizer: 'adam',
metrics: ['accuracy']
})
await model.fit(xTrain, yTrain, {
epochs: epochs,
callbacks: printCallback,
batchSize: batchSize,
validationData: [xTest, yTest]
//validationSplit: 0.2
})
const modelDirPath = path.join(__dirname, 'mymodel')
await model.save(`file://${modelDirPath}`)
}
const batchSize = 128
const epochs = 15
const numClasses = 10
const inputShape = [28, 28, 1]
const model = createModel(numClasses, inputShape)
//model.summary()
const dataset = loadData()
train(
model,
dataset,
batchSize,
epochs,
{ onEpochEnd: (epoch, log)=> { console.log(epoch, log) } }
)
node train.js してトレーニングを実行します。 ここでの環境（M1 macbook air）では 10分程度でトレーニングが完了しました。
保存したモデル ./mymodel/ が生成されていることを確認します。
推測する
train.js で生成したモデルを使って、実際に手描き数字を推測させてみます。
このエントリーでやったこととだいたい同じなので、 ここでは簡単な説明にとどめます。
自分で新たに書いた手描き数字データを ./images/ 以下に配置する.
0.jpg
1.jpg
2.jpg
3.jpg
4.jpg
5.jpg
6.jpg
7.jpg
8.jpg
9.jpg
そして推測するための infer.js は以下の通り。
const fs = require('fs')
const path = require('path')
const tf = require('@tensorflow/tfjs-node')
const range = (v)=>{ return [...Array(v).keys()] }
const infer = async ()=> {
const modelDirPath = path.join(__dirname, 'mymodel')
const modelFilePath = path.join(modelDirPath, 'model.json')
const model = await tf.loadLayersModel( `file://${modelFilePath}` )
const imagesDir = 'images'
const imageFilenames = fs.readdirSync(imagesDir)
const inputTensorList = imageFilenames.map((imageFilename)=> {
const image = fs.readFileSync(path.join(imagesDir, imageFilename))
const tensor = tf.node.decodeImage(image)
return tensor.div(255)
})
const inputTensors = tf.stack(inputTensorList)
console.log(inputTensors.shape)
const outputTensors = model.predict( inputTensors )
console.log( outputTensors.shape ) // [ 10, 10 ]
const maxValueIndexes = tf.argMax(outputTensors, 1).arraySync()
const maxValues = tf.max(outputTensors, 1).arraySync()
range(10).forEach((i)=> {
console.log( `- ${imageFilenames[i]} => ${maxValueIndexes[i]} (${maxValues[i]})` )
})
}
infer()
推測を実行。
$ node infer.js
[ 10, 28, 28, 1 ]
[ 10, 10 ]
- gray_0.jpg => 0 (0.9917171597480774)
- gray_1.jpg => 1 (0.9995224475860596)
- gray_2.jpg => 2 (0.9992350339889526)
- gray_3.jpg => 3 (0.9997997879981995)
- gray_4.jpg => 4 (0.9992501139640808)
- gray_5.jpg => 5 (0.9999731779098511)
- gray_6.jpg => 6 (0.9999617338180542)
- gray_7.jpg => 7 (0.9987815022468567)
- gray_8.jpg => 8 (0.9967820644378662)
- gray_9.jpg => 4 (0.9546715021133423)
9 以外はうまく推測できました。
以上です。