import os import json from loguru import logger import tensorflow as tf from ..dataset.batched_iterator import batched_iterator from ..io.find_paramsjson import find_paramsjson from ..io.readfile import readfile from ..io.writefile import writefile from .model_rainfallwater_contrastive import model_rainfallwater_contrastive from .helpers import make_callbacks from .helpers import summarywriter from .components.LayerContrastiveEncoder import LayerContrastiveEncoder from .components.LayerConvNeXtGamma import LayerConvNeXtGamma from .components.LayerCheeseMultipleOut import LayerCheeseMultipleOut from .helpers.summarywriter import summarywriter class RainfallWaterContraster(object): def __init__(self, dir_output=None, filepath_checkpoint=None, epochs=50, batch_size=64, **kwargs): super(RainfallWaterContraster, self).__init__() self.dir_output = dir_output self.epochs = epochs self.kwargs = kwargs self.batch_size = batch_size if filepath_checkpoint == None: if self.dir_output == None: raise Exception("Error: dir_output was not specified, and since no checkpoint was loaded training mode is activated.") if not os.path.exists(self.dir_output): os.mkdir(self.dir_output) self.filepath_summary = os.path.join(self.dir_output, "summary.txt") writefile(self.filepath_summary, "") # Empty the file ahead of time self.make_model() summarywriter(self.model, self.filepath_summary, append=True) writefile(os.path.join(self.dir_output, "params.json"), json.dumps(self.get_config())) else: self.load_model(filepath_checkpoint) def get_config(self): return { "epochs": self.epochs, "batch_size": self.batch_size, **self.kwargs } @staticmethod def from_checkpoint(filepath_checkpoint, **hyperparams): logger.info(f"Loading from checkpoint: {filepath_checkpoint}") return RainfallWaterContraster(filepath_checkpoint=filepath_checkpoint, **hyperparams) def make_model(self): self.model, self.model_predict = model_rainfallwater_contrastive( batch_size=self.batch_size, summary_file=self.filepath_summary, **self.kwargs ) def load_model(self, filepath_checkpoint): """ Loads a saved model from the given filename. filepath_checkpoint (string): The filepath to load the saved model from. """ self.model_predict = tf.keras.models.load_model(filepath_checkpoint, custom_objects={ "LayerContrastiveEncoder": LayerContrastiveEncoder, "LayerConvNeXtGamma": LayerConvNeXtGamma, "LayerCheeseMultipleOut": LayerCheeseMultipleOut }) def train(self, dataset_train, dataset_validate): return self.model.fit( dataset_train, validation_data=dataset_validate, epochs=self.epochs, callbacks=make_callbacks(self.dir_output, self.model_predict), steps_per_epoch=10 # For testing ) def embed(self, dataset): i_batch = -1 for batch in batched_iterator(dataset, tensors_in_item=2, batch_size=self.batch_size): i_batch += 1 rainfall = self.model_predict(batch[0], training=False) # ((rainfall, water), dummy_label) rainfall = tf.unstack(rainfall, axis=0) water = tf.unstack(batch[1], axis=0) for step_rainfall, step_water in zip(rainfall, water): yield step_rainfall, step_water # def embed_rainfall(self, dataset): # result = [] # for batch in dataset: # result_batch = self.model_predict(batch) # result.extend(tf.unstack(result_batch, axis=0)) # return result