diff --git a/aimodel/src/encoderonly_test_rainfall.py b/aimodel/src/encoderonly_test_rainfall.py index 5f101ee..6fcd575 100644 --- a/aimodel/src/encoderonly_test_rainfall.py +++ b/aimodel/src/encoderonly_test_rainfall.py @@ -4,8 +4,9 @@ import os import tensorflow as tf +from lib.dataset.dataset_encoderonly import dataset_encoderonly from lib.ai.components.convnext import make_convnext - +from lib.ai.helpers.summarywriter import summarywriter # ███████ ███ ██ ██ ██ # ██ ████ ██ ██ ██ @@ -14,8 +15,27 @@ from lib.ai.components.convnext import make_convnext # ███████ ██ ████ ████ # TODO: env vars & settings here +DIRPATH_INPUT = os.environ["DIRPATH_INPUT"] +DIRPATH_OUTPUT = os.environ["DIRPATH_OUTPUT"] +PATH_HEIGHTMAP = os.environ["PATH_HEIGHTMAP"] if "PATH_HEIGHTMAP" in os.environ else None +CHANNELS = os.environ["CHANNELS"] if "CHANNELS" in os.environ else 8 + +BATCH_SIZE = int(os.environ["BATCH_SIZE"]) if "BATCH_SIZE" in os.environ else 64 +WINDOW_SIZE = int(os.environ["WINDOW_SIZE"]) if "WINDOW_SIZE" in os.environ else 33 +STEPS_PER_EPOCH = int(os.environ["STEPS_PER_EPOCH"]) if "STEPS_PER_EPOCH" in os.environ else None + +logger.info("Encoder-only rainfall radar TEST") +logger.info(f"> DIRPATH_INPUT {DIRPATH_INPUT}") +logger.info(f"> DIRPATH_OUTPUT {DIRPATH_OUTPUT}") +logger.info(f"> PATH_HEIGHTMAP {PATH_HEIGHTMAP}") +logger.info(f"> CHANNELS {CHANNELS}") +logger.info(f"> BATCH_SIZE {BATCH_SIZE}") +logger.info(f"> WINDOW_SIZE {WINDOW_SIZE}") +logger.info(f"> STEPS_PER_EPOCH {STEPS_PER_EPOCH}") +if not os.path.exists(DIRPATH_OUTPUT): + os.makedirs(os.path.join(DIRPATH_OUTPUT, "checkpoints")) # ██████ █████ ████████ █████ ███████ ███████ ████████ @@ -24,7 +44,12 @@ from lib.ai.components.convnext import make_convnext # ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ # ██████ ██ ██ ██ ██ ██ ███████ ███████ ██ - +dataset_train, dataset_validate = dataset_encoderonly( + dirpath_input=DIRPATH_INPUT, + filepath_heightmap=PATH_HEIGHTMAP, + batch_size=BATCH_SIZE, + windowsize=WINDOW_SIZE +) # ███ ███ ██████ ██████ ███████ ██ @@ -33,9 +58,42 @@ from lib.ai.components.convnext import make_convnext # ██ ██ ██ ██ ██ ██ ██ ██ ██ # ██ ██ ██████ ██████ ███████ ███████ +def make_encoderonly(windowsize, channels, encoder="convnext", water_bins=2): + if encoder == "convnext": + model = make_convnext(input_shape=(windowsize, windowsize, channels), num_classes=water_bins, **kwargs) + elif encoder == "resnet": + layer_in = tf.keras.Input(shape=(windowsize, windowsize, channels)) + layer_next = tf.keras.applications.resnet50.ResNet50( + weights=None, + include_top=True, + classes=water_bins, + input_tensor=layer_in, + pooling="max", + ) + + model = tf.keras.Model( + inputs = layer_in, + outputs = layer_next + ) + else: + raise Exception(f"Error: Unknown encoder '{encoder}' (known encoders: convnext, resnet).") + + model.compile( + optimizer="Adam", + loss = tf.keras.losses.SparseCategoricalCrossentropy(), + metrics = [ + tf.keras.metrics.SparseCategoricalAccuracy() + ] + ) + + return model -# TODO: Have an +model = make_encoderonly( + windowsize=WINDOW_SIZE, + channels=CHANNELS +) +summarywriter(model, os.path.join(DIRPATH_OUTPUT, "summary.txt")) # ████████ ██████ █████ ██ ███ ██ ██ ███ ██ ██████ # ██ ██ ██ ██ ██ ██ ████ ██ ██ ████ ██ ██ @@ -43,6 +101,28 @@ from lib.ai.components.convnext import make_convnext # ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ ██ # ██ ██ ██ ██ ██ ██ ██ ████ ██ ██ ████ ██████ +model.fit(dataset_train, + validation_data=dataset_validate, + epochs=25, + + callbacks=[ + tf.keras.callbacks.CSVLogger( + filename=os.path.join(DIRPATH_OUTPUT, "metrics.tsv"), + separator="\t" + ), + CallbackCustomModelCheckpoint( + model_to_checkpoint=model, + filepath=os.path.join( + DIRPATH_OUTPUT, + "checkpoints" + "checkpoint_e{epoch:d}_loss{loss:.3f}.hdf5" + ), + monitor="loss" + ), + ], + steps_per_epoch=STEPS_PER_EPOCH, +) +logger.info(">>> Training complete") # ██████ ██████ ███████ ██████ ██ ██████ ████████ ██ ██████ ███ ██ @@ -52,7 +132,7 @@ from lib.ai.components.convnext import make_convnext # ██ ██ ██ ███████ ██████ ██ ██████ ██ ██ ██████ ██ ████ - +# TODO FILL THIS IN diff --git a/aimodel/src/lib/dataset/dataset_encoderonly.py b/aimodel/src/lib/dataset/dataset_encoderonly.py index b9ccb25..2a6ce13 100644 --- a/aimodel/src/lib/dataset/dataset_encoderonly.py +++ b/aimodel/src/lib/dataset/dataset_encoderonly.py @@ -15,19 +15,15 @@ from .parse_heightmap import parse_heightmap # TO PARSE: -def parse_item(metadata, output_size=100, input_size="same", water_threshold=0.1, water_bins=2, heightmap=None, rainfall_scale_up=1): +def parse_item(metadata, water_threshold=0.1, water_bins=2, heightmap=None, rainfall_scale_up=1, windowsize=33): if input_size == "same": input_size = output_size # This is almost always the case with e.g. the DeepLabV3+ model water_height_source, water_width_source = metadata["waterdepth"] - water_offset_x = math.ceil((water_width_source - output_size) / 2) - water_offset_y = math.ceil((water_height_source - output_size) / 2) rainfall_channels, rainfall_height_source, rainfall_width_source = metadata["rainfallradar"] rainfall_height_source *= rainfall_scale_up rainfall_width_source *= rainfall_scale_up - rainfall_offset_x = math.ceil((rainfall_width_source - input_size) / 2) - rainfall_offset_y = math.ceil((rainfall_height_source - input_size) / 2) print("DEBUG DATASET:rainfall shape", metadata["rainfallradar"], "/", f"w {rainfall_width_source} h {rainfall_height_source}") print("DEBUG DATASET:water shape", metadata["waterdepth"]) @@ -35,8 +31,6 @@ def parse_item(metadata, output_size=100, input_size="same", water_threshold=0.1 print("DEBUG DATASET:water_bins", water_bins) print("DEBUG DATASET:output_size", output_size) print("DEBUG DATASET:input_size", input_size) - print("DEBUG DATASET:water_offset x", water_offset_x, "y", water_offset_y) - print("DEBUG DATASET:rainfall_offset x", rainfall_offset_x, "y", rainfall_offset_y) if heightmap is not None: heightmap = tf.expand_dims(heightmap, axis=-1) @@ -72,21 +66,14 @@ def parse_item(metadata, output_size=100, input_size="same", water_threshold=0.1 rainfall = tf.concat([rainfall, heightmap], axis=-1) if rainfall_scale_up > 1: rainfall = tf.repeat(tf.repeat(rainfall, rainfall_scale_up, axis=0), rainfall_scale_up, axis=1) - if input_size is not None: - rainfall = tf.image.crop_to_bounding_box(rainfall, - offset_width=rainfall_offset_x, - offset_height=rainfall_offset_y, - target_width=input_size, - target_height=input_size, - ) # rainfall = tf.image.resize(rainfall, tf.cast(tf.constant(metadata["rainfallradar"]) / 2, dtype=tf.int32)) water = tf.expand_dims(water, axis=-1) # [height, width] → [height, width, channels=1] - water = tf.image.crop_to_bounding_box(water, - offset_width=water_offset_x, - offset_height=water_offset_y, - target_width=output_size, - target_height=output_size + water = tf.image.crop_to_bounding_box(water, # for path generation; we only want the water bit where the patches convolve + offset_width=math.floor(windowsize/2), + offset_height=math.floor(windowsize/2), + target_width=water_width_source - (math.floor(windowsize/2)*2), + target_height=water_height_source - (math.floor(windowsize/2)*2) ) print("DEBUG:dataset BEFORE_SQUEEZE water", water.shape) @@ -96,7 +83,7 @@ def parse_item(metadata, output_size=100, input_size="same", water_threshold=0.1 # water = tf.cast(tf.math.greater_equal(water, water_threshold), dtype=tf.int32) # water = tf.one_hot(water, water_bins, axis=-1, dtype=tf.int32) # LOSS dice - water = tf.cast(tf.math.greater_equal(water, water_threshold), dtype=tf.float32) + water = tf.cast(tf.math.greater_equal(water, water_threshold), dtype=tf.int32) rainfall = tf.image.extract_patches(tf.expand_dims(rainfall, axis=0), sizes=[1,windowsize,windowsize], @@ -106,7 +93,7 @@ def parse_item(metadata, output_size=100, input_size="same", water_threshold=0.1 ) rainfall = tf.reshape(rainfall, [-1, windowsize, windowsize, rainfall_channels]) - # TODO: extract single water values here to match the above rainfall patches + water = tf.reshape([-1]) # we flatten because we cropped to the right shape above print("DEBUG DATASET_OUT:rainfall shape", rainfall.shape) print("DEBUG DATASET_OUT:water shape", water.shape) @@ -114,7 +101,7 @@ def parse_item(metadata, output_size=100, input_size="same", water_threshold=0.1 return tf.function(parse_item_inner) -def make_dataset(filepaths, compression_type="GZIP", parallel_reads_multiplier=1.5, shuffle_buffer_size=128, batch_size=64, prefetch=True, shuffle=True, filepath_heightmap=None, **kwargs): +def make_dataset(filepaths, compression_type="GZIP", parallel_reads_multiplier=3, shuffle_buffer_size=2**16, batch_size=64, prefetch=True, shuffle=True, filepath_heightmap=None, **kwargs): if "NO_PREFETCH" in os.environ: logger.info("disabling data prefetching.") @@ -128,10 +115,16 @@ def make_dataset(filepaths, compression_type="GZIP", parallel_reads_multiplier=1 num_parallel_reads=math.ceil(os.cpu_count() * parallel_reads_multiplier) if parallel_reads_multiplier > 0 else None ) if shuffle: - dataset = dataset.shuffle(shuffle_buffer_size) + dataset = dataset.shuffle(128) # additional shuffle buffer to mix things up + dataset = dataset.map(parse_item(heightmap=heightmap, **kwargs), num_parallel_calls=tf.data.AUTOTUNE) \ .unbatch() + if shuffle: + # memory used = (windowsize*windowsize + 1) * buffersize * channels + # defaults = (33*33 + 1) * 2**16 * 8 = about 2.219GiB + dataset = dataset.shuffle(shuffle_buffer_size) + if batch_size != None: dataset = dataset.batch(batch_size, drop_remainder=True) if prefetch: @@ -152,7 +145,7 @@ def get_filepaths(dirpath_input, do_shuffle=True): return result -def dataset_mono(dirpath_input, train_percentage=0.8, **kwargs): +def dataset_encoderonly(dirpath_input, train_percentage=0.8, **kwargs): filepaths = get_filepaths(dirpath_input) filepaths_count = len(filepaths) dataset_splitpoint = math.floor(filepaths_count * train_percentage) @@ -167,7 +160,7 @@ def dataset_mono(dirpath_input, train_percentage=0.8, **kwargs): return dataset_train, dataset_validate #, filepaths -def dataset_mono_predict(dirpath_input, **kwargs): +def dataset_encoderonly_predict(dirpath_input, **kwargs): """Creates a tf.data.Dataset() for prediction using the contrastive learning model. Note that this WILL MANGLE THE ORDERING if you set parallel_reads_multiplier to anything other than 0!!