|
import tensorflow as tf |
|
import tensorflow_datasets as tfds |
|
from tensorflow_examples.models.pix2pix import pix2pix |
|
import helper |
|
|
|
|
|
dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True) |
|
|
|
TRAIN_LENGTH = info.splits['train'].num_examples |
|
BATCH_SIZE = 64 |
|
BUFFER_SIZE = 1000 |
|
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE |
|
|
|
train_images = dataset['train'].map(helper.load_image, num_parallel_calls=tf.data.AUTOTUNE) |
|
test_images = dataset['test'].map(helper.oad_image, num_parallel_calls=tf.data.AUTOTUNE) |
|
|
|
train_batches = (train_images.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat().map(helper.Augment()).prefetch(buffer_size=tf.data.AUTOTUNE)) |
|
test_batches = test_images.batch(BATCH_SIZE) |
|
|
|
|
|
base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False) |
|
|
|
layer_names = [ |
|
'block_1_expand_relu', |
|
'block_3_expand_relu', |
|
'block_6_expand_relu', |
|
'block_13_expand_relu', |
|
'block_16_project', |
|
] |
|
base_model_outputs = [base_model.get_layer(name).output for name in layer_names] |
|
|
|
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs) |
|
|
|
down_stack.trainable = False |
|
|
|
up_stack = [pix2pix.upsample(512, 3), pix2pix.upsample(256, 3), pix2pix.upsample(128, 3), pix2pix.upsample(64, 3),] |
|
|
|
OUTPUT_CLASSES = 3 |
|
|
|
model = helper.U_net_model(OUTPUT_CLASSES, down_stack, up_stack) |
|
model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy']) |
|
|
|
EPOCHS = 20 |
|
VAL_SUBSPLITS = 5 |
|
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS |
|
|
|
model.fit(train_batches, epochs=EPOCHS, steps_per_epoch=STEPS_PER_EPOCH, validation_steps=VALIDATION_STEPS, validation_data=test_batches) |
|
|
|
model.save("pets.h5") |
|
model.save("pets.keras") |
|
model.save("model/dogs") |
|
|