add optional PATH_CHECKPOINT env var

This commit is contained in:
Starbeamrainbowlabs 2023-01-11 17:20:19 +00:00
parent 0e3de8f5fc
commit 93e663e45d
Signed by: sbrl
GPG key ID: 1BE5172E637709C2

View file

@ -28,6 +28,8 @@ STEPS_PER_EPOCH = int(os.environ["STEPS_PER_EPOCH"]) if "STEPS_PER_EPOCH" in os.
DIR_OUTPUT=os.environ["DIR_OUTPUT"] if "DIR_OUTPUT" in os.environ else f"output/{datetime.utcnow().date().isoformat()}_deeplabv3plus_rainfall_TEST" DIR_OUTPUT=os.environ["DIR_OUTPUT"] if "DIR_OUTPUT" in os.environ else f"output/{datetime.utcnow().date().isoformat()}_deeplabv3plus_rainfall_TEST"
PATH_CHECKPOINT = os.environ["PATH_CHECKPOINT"] if "PATH_CHECKPOINT" in os.environ else None
if not os.path.exists(DIR_OUTPUT): if not os.path.exists(DIR_OUTPUT):
os.makedirs(os.path.join(DIR_OUTPUT, "checkpoints")) os.makedirs(os.path.join(DIR_OUTPUT, "checkpoints"))
@ -38,6 +40,7 @@ logger.info(f"> PATH_HEIGHTMAP {PATH_HEIGHTMAP}")
logger.info(f"> PATH_COLOURMAP {PATH_COLOURMAP}") logger.info(f"> PATH_COLOURMAP {PATH_COLOURMAP}")
logger.info(f"> STEPS_PER_EPOCH {STEPS_PER_EPOCH}") logger.info(f"> STEPS_PER_EPOCH {STEPS_PER_EPOCH}")
logger.info(f"> DIR_OUTPUT {DIR_OUTPUT}") logger.info(f"> DIR_OUTPUT {DIR_OUTPUT}")
logger.info(f"> PATH_CHECKPOINT {PATH_CHECKPOINT}")
dataset_train, dataset_validate = dataset_mono( dataset_train, dataset_validate = dataset_mono(
@ -60,14 +63,16 @@ logger.info("Validation Dataset:", dataset_validate)
# ██ ██ ██ ██ ██ ██ ██ ██ ██ # ██ ██ ██ ██ ██ ██ ██ ██ ██
# ██ ██ ██████ ██████ ███████ ███████ # ██ ██ ██████ ██████ ███████ ███████
def convolution_block(
if PATH_CHECKPOINT is None:
def convolution_block(
block_input, block_input,
num_filters=256, num_filters=256,
kernel_size=3, kernel_size=3,
dilation_rate=1, dilation_rate=1,
padding="same", padding="same",
use_bias=False, use_bias=False,
): ):
x = tf.keras.layers.Conv2D( x = tf.keras.layers.Conv2D(
num_filters, num_filters,
kernel_size=kernel_size, kernel_size=kernel_size,
@ -80,7 +85,7 @@ def convolution_block(
return tf.nn.relu(x) return tf.nn.relu(x)
def DilatedSpatialPyramidPooling(dspp_input): def DilatedSpatialPyramidPooling(dspp_input):
dims = dspp_input.shape dims = dspp_input.shape
x = tf.keras.layers.AveragePooling2D(pool_size=(dims[-3], dims[-2]))(dspp_input) x = tf.keras.layers.AveragePooling2D(pool_size=(dims[-3], dims[-2]))(dspp_input)
x = convolution_block(x, kernel_size=1, use_bias=True) x = convolution_block(x, kernel_size=1, use_bias=True)
@ -98,7 +103,7 @@ def DilatedSpatialPyramidPooling(dspp_input):
return output return output
def DeeplabV3Plus(image_size, num_classes, num_channels=3): def DeeplabV3Plus(image_size, num_classes, num_channels=3):
model_input = tf.keras.Input(shape=(image_size, image_size, num_channels)) model_input = tf.keras.Input(shape=(image_size, image_size, num_channels))
resnet50 = tf.keras.applications.ResNet50( resnet50 = tf.keras.applications.ResNet50(
weights="imagenet" if num_channels == 3 else None, weights="imagenet" if num_channels == 3 else None,
@ -125,8 +130,8 @@ def DeeplabV3Plus(image_size, num_classes, num_channels=3):
return tf.keras.Model(inputs=model_input, outputs=model_output) return tf.keras.Model(inputs=model_input, outputs=model_output)
model = DeeplabV3Plus(image_size=IMAGE_SIZE, num_classes=NUM_CLASSES, num_channels=8) model = DeeplabV3Plus(image_size=IMAGE_SIZE, num_classes=NUM_CLASSES, num_channels=8)
summarywriter(model, os.path.join(DIR_OUTPUT, "summary.txt")) summarywriter(model, os.path.join(DIR_OUTPUT, "summary.txt"))
@ -137,14 +142,15 @@ summarywriter(model, os.path.join(DIR_OUTPUT, "summary.txt"))
# ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ # ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██
# ██ ██ ██ ██ ██ ██ ██ ████ ██ ██ ████ ██████ # ██ ██ ██ ██ ██ ██ ██ ████ ██ ██ ████ ██████
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) if PATH_CHECKPOINT is None:
model.compile( loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
loss=loss, loss=loss,
metrics=["accuracy"], metrics=["accuracy"],
) )
logger.info(">>> Beginning training") logger.info(">>> Beginning training")
history = model.fit(dataset_train, history = model.fit(dataset_train,
validation_data=dataset_validate, validation_data=dataset_validate,
epochs=25, epochs=25,
callbacks=[ callbacks=[
@ -163,37 +169,37 @@ history = model.fit(dataset_train,
), ),
], ],
steps_per_epoch=STEPS_PER_EPOCH, steps_per_epoch=STEPS_PER_EPOCH,
) )
logger.info(">>> Training complete") logger.info(">>> Training complete")
logger.info(">>> Plotting graphs") logger.info(">>> Plotting graphs")
plt.plot(history.history["loss"]) plt.plot(history.history["loss"])
plt.title("Training Loss") plt.title("Training Loss")
plt.ylabel("loss") plt.ylabel("loss")
plt.xlabel("epoch") plt.xlabel("epoch")
plt.savefig(os.path.join(DIR_OUTPUT, "loss.png")) plt.savefig(os.path.join(DIR_OUTPUT, "loss.png"))
plt.close() plt.close()
plt.plot(history.history["accuracy"]) plt.plot(history.history["accuracy"])
plt.title("Training Accuracy") plt.title("Training Accuracy")
plt.ylabel("accuracy") plt.ylabel("accuracy")
plt.xlabel("epoch") plt.xlabel("epoch")
plt.savefig(os.path.join(DIR_OUTPUT, "acc.png")) plt.savefig(os.path.join(DIR_OUTPUT, "acc.png"))
plt.close() plt.close()
plt.plot(history.history["val_loss"]) plt.plot(history.history["val_loss"])
plt.title("Validation Loss") plt.title("Validation Loss")
plt.ylabel("val_loss") plt.ylabel("val_loss")
plt.xlabel("epoch") plt.xlabel("epoch")
plt.savefig(os.path.join(DIR_OUTPUT, "val_loss.png")) plt.savefig(os.path.join(DIR_OUTPUT, "val_loss.png"))
plt.close() plt.close()
plt.plot(history.history["val_accuracy"]) plt.plot(history.history["val_accuracy"])
plt.title("Validation Accuracy") plt.title("Validation Accuracy")
plt.ylabel("val_accuracy") plt.ylabel("val_accuracy")
plt.xlabel("epoch") plt.xlabel("epoch")
plt.savefig(os.path.join(DIR_OUTPUT, "val_acc.png")) plt.savefig(os.path.join(DIR_OUTPUT, "val_acc.png"))
plt.close() plt.close()