class UNet(BaseModel): """Unet Model Class""" def __init__(self, config): super().__init__(config) # self.base_model = tf.keras.applications.MobileNetV2(input_shape=self.config.model.input, include_top=False) #https: // www.tensorflow.org / api_docs / python / tf / keras / applications?hl = de self.model = None self.output_channels = self.config.model.output self.dataset = None self.info = None self.batch_size = self.config.train.batch_size #self.buffer_size = self.config.train.buffer_size self.epoches = self.config.train.epoches #self.val_subsplits = self.config.train.val_subsplits self.validation_steps = 0 self.train_length = 0 self.steps_per_epoch = 0 self.image_size = self.config.data.image_size self.train_dataset = [] self.test_dataset = [] def load_data(self): """Loads and Preprocess data """ self.dataset = DataLoader().load_data(self.config) self._preprocess_data() def _preprocess_data(self): """ Splits into training and test and set training parameters""" self.train_dataset = self.dataset['generated_training_batches'] self.test_dataset = self.dataset['generated_test_batches'] #self._set_training_parameters() # def _set_training_parameters(self): # """Sets training parameters""" #self.train_length = self.info.splits['train'].num_examples #self.steps_per_epoch = self.train_length // self.batch_size #self.validation_steps = self.info.splits['test'].num_examples // self.batch_size // self.val_subsplits # def _normalize(self, input_image, input_mask): # !!! normalization by scaling happends in the dataloader.py by ImageDataGenerator !!! # """ Normalise input image # Args: # input_image (tf.image): The input image # input_mask (int): The image mask # # Returns: # input_image (tf.image): The normalized input image # input_mask (int): The new image mask # """ # input_image = tf.cast(input_image, tf.float32) / 255.0 # input_mask -= 1 # return input_image, input_mask # @tf.function # def _load_image_train(self, datapoint): # """ Loads and preprocess a single training image """ # input_image = tf.image.resize(datapoint['image'], (self.image_size, self.image_size)) # input_mask = tf.image.resize(datapoint['segmentation_mask'], (self.image_size, self.image_size)) # # if tf.random.uniform(()) > 0.5: # input_image = tf.image.flip_left_right(input_image) # input_mask = tf.image.flip_left_right(input_mask) # # input_image, input_mask = self._normalize(input_image, input_mask) # # return input_image, input_mask # def _load_image_test(self, datapoint): # """ Loads and preprocess a single test images""" # # input_image = tf.image.resize(datapoint['image'], (self.image_size, self.image_size)) # input_mask = tf.image.resize(datapoint['segmentation_mask'], (self.image_size, self.image_size)) # # input_image, input_mask = self._normalize(input_image, input_mask) # # return input_image, input_mask def build(self): """ Builds the Keras model based """ layer_names = [ 'block_1_expand_relu', # 64x64 'block_3_expand_relu', # 32x32 'block_6_expand_relu', # 16x16 'block_13_expand_relu', # 8x8 'block_16_project', # 4x4 ] layers = [self.base_model.get_layer(name).output for name in layer_names] # Create the feature extraction model # down_stack = tf.keras.Model(inputs=self.base_model.input, outputs=layers) # down_stack.trainable = False up_stack = [ pix2pix.upsample(self.config.model.up_stack.layer_1, self.config.model.up_stack.kernels), # 4x4 -> 8x8 pix2pix.upsample(self.config.model.up_stack.layer_2, self.config.model.up_stack.kernels), # 8x8 -> 16x16 pix2pix.upsample(self.config.model.up_stack.layer_3, self.config.model.up_stack.kernels), # 16x16 -> 32x32 pix2pix.upsample(self.config.model.up_stack.layer_4, self.config.model.up_stack.kernels), # 32x32 -> 64x64 ] inputs = tf.keras.layers.Input(shape=self.config.model.input) x = inputs # Downsampling through the model # skips = down_stack(x) # x = skips[-1] # skips = reversed(skips[:-1]) # Upsampling and establishing the skip connections # for up, skip in zip(up_stack, skips): # x = up(x) # concat = tf.keras.layers.Concatenate() # x = concat([x, skip]) # This is the last layer of the model last = tf.keras.layers.Conv2DTranspose( self.output_channels, self.config.model.up_stack.kernels, strides=2, padding='same') # 64x64 -> 128x128 x = last(x) self.model = tf.keras.Model(inputs=inputs, outputs=x) def train(self): """Compiles and trains the model""" self.model.compile(optimizer=self.config.train.optimizer.type, loss=tf.keras.losses.BinaryCrossentropy(), metrics=self.config.train.metrics) early_stop = EarlyStopping(monitor=self.config.train.EarlyStopping.monitor, mode=self.config.train.EarlyStopping.mode, verbose=1, patience=self.config.train.EarlyStopping.self.config.train.EarlyStopping.patience) model_history = self.model.fit_generator(self.train_dataset, epochs=self.epoches, steps_per_epoch=self.steps_per_epoch, validation_steps=self.validation_steps, callbacks=[early_stop], validation_data=self.test_dataset) return model_history.history.history def evaluate(self): """Predicts resuts for the test dataset""" predictions = [] for image, mask in self.dataset.take(1): predictions.append(self.model.predict(image)) return predictions def distributed_train(self): mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"]) with mirrored_strategy.scope(): self.model = tf.keras.Model(inputs=inputs, outputs=x) self.model.compile(...) self.model.fit(...) os.environ["TF_CONFIG"] = json.dumps( { "cluster":{ "worker": ["host1:port", "host2:port", "host3:port"] }, "task":{ "type": "worker", "index": 1 } } ) multi_worker_mirrored_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy() with multi_worker_mirrored_strategy.scope(): self.model = tf.keras.Model(inputs=inputs, outputs=x) self.model.compile(...) self.model.fit(...) parameter_server_strategy = tf.distribute.experimental.ParameterServerStrategy() os.environ["TF_CONFIG"] = json.dumps( { "cluster": { "worker": ["host1:port", "host2:port", "host3:port"], "ps": ["host4:port", "host5:port"] }, "task": { "type": "worker", "index": 1 } } )
class UNet(BaseModel): """Unet Model Class""" def __init__(self, config): super().__init__(config) self.base_model = tf.keras.applications.MobileNetV2(input_shape=self.config.model.input, include_top=False) self.model = None self.output_channels = self.config.model.output self.dataset = None self.info = None self.batch_size = self.config.train.batch_size self.buffer_size = self.config.train.buffer_size self.epoches = self.config.train.epoches self.val_subsplits = self.config.train.val_subsplits self.validation_steps = 0 self.train_length = 0 self.steps_per_epoch = 0 self.image_size = self.config.data.image_size self.train_dataset = [] self.test_dataset = [] def load_data(self): """Loads and Preprocess data """ self.dataset, self.info = DataLoader().load_data(self.config.data) self._preprocess_data() def _preprocess_data(self): """ Splits into training and test and set training parameters""" train = self.dataset['train'].map(self._load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE) test = self.dataset['test'].map(self._load_image_test) self.train_dataset = train.cache().shuffle(self.buffer_size).batch(self.batch_size).repeat() self.train_dataset = self.train_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) self.test_dataset = test.batch(self.batch_size) self._set_training_parameters() def _set_training_parameters(self): """Sets training parameters""" self.train_length = self.info.splits['train'].num_examples self.steps_per_epoch = self.train_length // self.batch_size self.validation_steps = self.info.splits['test'].num_examples // self.batch_size // self.val_subsplits def _normalize(self, input_image, input_mask): """ Normalise input image Args: input_image (tf.image): The input image input_mask (int): The image mask Returns: input_image (tf.image): The normalized input image input_mask (int): The new image mask """ input_image = tf.cast(input_image, tf.float32) / 255.0 input_mask -= 1 return input_image, input_mask @tf.function def _load_image_train(self, datapoint): """ Loads and preprocess a single training image """ input_image = tf.image.resize(datapoint['image'], (self.image_size, self.image_size)) input_mask = tf.image.resize(datapoint['segmentation_mask'], (self.image_size, self.image_size)) if tf.random.uniform(()) > 0.5: input_image = tf.image.flip_left_right(input_image) input_mask = tf.image.flip_left_right(input_mask) input_image, input_mask = self._normalize(input_image, input_mask) return input_image, input_mask def _load_image_test(self, datapoint): """ Loads and preprocess a single test images""" input_image = tf.image.resize(datapoint['image'], (self.image_size, self.image_size)) input_mask = tf.image.resize(datapoint['segmentation_mask'], (self.image_size, self.image_size)) input_image, input_mask = self._normalize(input_image, input_mask) return input_image, input_mask def build(self): """ Builds the Keras model based """ layer_names = [ 'block_1_expand_relu', # 64x64 'block_3_expand_relu', # 32x32 'block_6_expand_relu', # 16x16 'block_13_expand_relu', # 8x8 'block_16_project', # 4x4 ] layers = [self.base_model.get_layer(name).output for name in layer_names] # Create the feature extraction model down_stack = tf.keras.Model(inputs=self.base_model.input, outputs=layers) down_stack.trainable = False up_stack = [ pix2pix.upsample(self.config.model.up_stack.layer_1, self.config.model.up_stack.kernels), # 4x4 -> 8x8 pix2pix.upsample(self.config.model.up_stack.layer_2, self.config.model.up_stack.kernels), # 8x8 -> 16x16 pix2pix.upsample(self.config.model.up_stack.layer_3, self.config.model.up_stack.kernels), # 16x16 -> 32x32 pix2pix.upsample(self.config.model.up_stack.layer_4, self.config.model.up_stack.kernels), # 32x32 -> 64x64 ] inputs = tf.keras.layers.Input(shape=self.config.model.input) x = inputs # Downsampling through the model skips = down_stack(x) x = skips[-1] skips = reversed(skips[:-1]) # Upsampling and establishing the skip connections for up, skip in zip(up_stack, skips): x = up(x) concat = tf.keras.layers.Concatenate() x = concat([x, skip]) # This is the last layer of the model last = tf.keras.layers.Conv2DTranspose( self.output_channels, self.config.model.up_stack.kernels, strides=2, padding='same') # 64x64 -> 128x128 x = last(x) self.model = tf.keras.Model(inputs=inputs, outputs=x) def train(self): """Compiles and trains the model""" self.model.compile(optimizer=self.config.train.optimizer.type, loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=self.config.train.metrics) model_history = self.model.fit(self.train_dataset, epochs=self.epoches, steps_per_epoch=self.steps_per_epoch, validation_steps=self.validation_steps, validation_data=self.test_dataset) return model_history.history['loss'], model_history.history['val_loss'] def evaluate(self): """Predicts resuts for the test dataset""" predictions = [] for image, mask in self.dataset.take(1): predictions.append(self.model.predict(image)) return predictions