48 lines
1.1 KiB
Python
48 lines
1.1 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='int',
|
|
seed=seed,
|
|
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='int',
|
|
seed=seed,
|
|
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="relu"),
|
|
{{- else if eq .LayerType 3}}
|
|
layers.Flatten(),
|
|
{{- else }}
|
|
ERROR
|
|
{{- end }}
|
|
{{- end }}
|
|
])
|
|
|
|
model.compile(loss=losses.SparseCategoricalCrossentropy(), optimizer=tf.keras.optimizers.Adam())
|
|
|
|
his = model.fit(dataset, validation_data= dataset_validation, epochs=100)
|
|
|