62 lines
1.3 KiB
Python
62 lines
1.3 KiB
Python
import tensorflow as tf
|
|
import random
|
|
from tensorflow import keras
|
|
from keras import layers, losses, optimizers
|
|
|
|
seed = random.randint(0, 100000000)
|
|
|
|
batch_size = 100
|
|
|
|
dataset = keras.utils.image_dataset_from_directory(
|
|
"{{ .DataDir }}",
|
|
color_mode="rgb",
|
|
validation_split=0.2,
|
|
label_mode='categorical',
|
|
seed=seed,
|
|
shuffle=True,
|
|
subset="training",
|
|
image_size=({{ .Size }}),
|
|
batch_size=batch_size)
|
|
|
|
dataset_validation = keras.utils.image_dataset_from_directory(
|
|
"{{ .DataDir }}",
|
|
color_mode="rgb",
|
|
validation_split=0.2,
|
|
label_mode='categorical',
|
|
seed=seed,
|
|
shuffle=True,
|
|
subset="validation",
|
|
image_size=({{ .Size }}),
|
|
batch_size=batch_size)
|
|
|
|
model = keras.Sequential([
|
|
{{- range .Layers }}
|
|
{{- if eq .LayerType 1}}
|
|
layers.Rescaling(1./255),
|
|
{{- else if eq .LayerType 2 }}
|
|
layers.Dense({{ .Shape }}, activation="sigmoid"),
|
|
{{- else if eq .LayerType 3}}
|
|
layers.Flatten(),
|
|
{{- else }}
|
|
ERROR
|
|
{{- end }}
|
|
{{- end }}
|
|
])
|
|
|
|
model.compile(
|
|
loss=losses.CategoricalCrossentropy(),
|
|
optimizer=tf.keras.optimizers.Adam(),
|
|
metrics=['accuracy'])
|
|
|
|
his = model.fit(dataset, validation_data= dataset_validation, epochs=50)
|
|
|
|
acc = his.history["accuracy"]
|
|
|
|
f = open("accuracy.val", "w")
|
|
f.write(str(acc[-1]))
|
|
f.close()
|
|
|
|
tf.saved_model.save(model, "model")
|
|
|
|
# model.save("model.keras", save_format="tf")
|