def __init__( self, checkpoint=None, text_postproc=None, data_preproc=None, codec=None, network=None, batch_size=1, processes=1, auto_update_checkpoints=True, with_gt=False, ctc_decoder_params=None, ): """ Predicting a dataset based on a trained model Parameters ---------- checkpoint : str, optional filepath of the checkpoint of the network to load, alternatively you can directly use a loaded `network` text_postproc : TextProcessor, optional text processor to be applied on the predicted sentence for the final output. If loaded from a checkpoint the text processor will be loaded from it. data_preproc : DataProcessor, optional data processor (must be the same as of the trained model) to be applied to the input image. If loaded from a checkpoint the text processor will be loaded from it. codec : Codec, optional Codec of the deep net to use for decoding. This parameter is only required if a custom codec is used, or a `network` has been provided instead of a `checkpoint` network : ModelInterface, optional DNN instance to used. Alternatively you can provide a `checkpoint` to load a network. batch_size : int, optional Batch size to use for prediction processes : int, optional The number of processes to use for prediction auto_update_checkpoints : bool, optional Update old models automatically (this will change the checkpoint files) with_gt : bool, optional The prediction will also output the ground truth if available else None ctc_decoder_params : optional Parameters of the ctc decoder """ self.network = network self.checkpoint = checkpoint self.processes = processes self.auto_update_checkpoints = auto_update_checkpoints self.with_gt = with_gt if checkpoint: if network: raise Exception( "Either a checkpoint or a network can be provided") ckpt = Checkpoint(checkpoint, auto_update=self.auto_update_checkpoints) self.checkpoint = ckpt.ckpt_path checkpoint_params = ckpt.checkpoint self.model_params = checkpoint_params.model self.codec = codec if codec else Codec( self.model_params.codec.charset) self.network_params = self.model_params.network backend = create_backend_from_checkpoint( checkpoint_params=checkpoint_params, processes=processes) self.text_postproc = text_postproc if text_postproc else text_processor_from_proto( self.model_params.text_postprocessor, "post") self.data_preproc = data_preproc if data_preproc else data_processor_from_proto( self.model_params.data_preprocessor) self.network = backend.create_net( codec=self.codec, ctc_decoder_params=ctc_decoder_params, checkpoint_to_load=ckpt, graph_type="predict", batch_size=batch_size) elif network: self.codec = codec self.model_params = None self.network_params = network.network_proto self.text_postproc = text_postproc self.data_preproc = data_preproc if not codec: raise Exception( "A codec is required if preloaded network is used.") else: raise Exception( "Either a checkpoint or a existing backend must be provided") self.out_to_in_trans = OutputToInputTransformer( self.data_preproc, self.network)
def train(self, auto_compute_codec=False, progress_bar=False, training_callback=ConsoleTrainingCallback()): """ Launch the training Parameters ---------- auto_compute_codec : bool Compute the codec automatically based on the provided ground truth. Else provide a codec using a whitelist (faster). progress_bar : bool Show or hide any progress bar training_callback : TrainingCallback Callback for the training process (e.g., for displaying the current cer, loss in the console) """ with ExitStackWithPop() as exit_stack: checkpoint_params = self.checkpoint_params train_start_time = time.time() + self.checkpoint_params.total_time exit_stack.enter_context(self.dataset) if self.validation_dataset: exit_stack.enter_context(self.validation_dataset) # load training dataset if self.preload_training: new_dataset = self.dataset.to_raw_input_dataset(processes=checkpoint_params.processes, progress_bar=progress_bar) exit_stack.pop(self.dataset) self.dataset = new_dataset exit_stack.enter_context(self.dataset) # load validation dataset if self.validation_dataset and self.preload_validation: new_dataset = self.validation_dataset.to_raw_input_dataset(processes=checkpoint_params.processes, progress_bar=progress_bar) exit_stack.pop(self.validation_dataset) self.validation_dataset = new_dataset exit_stack.enter_context(self.validation_dataset) # compute the codec if self.codec: codec = self.codec else: if len(self.codec_whitelist) == 0 or auto_compute_codec: codec = Codec.from_input_dataset([self.dataset, self.validation_dataset], whitelist=self.codec_whitelist, progress_bar=progress_bar) else: codec = Codec.from_texts([], whitelist=self.codec_whitelist) # create backend network_params = checkpoint_params.model.network network_params.features = checkpoint_params.model.line_height if self.weights: # if we load the weights, take care of codec changes as-well ckpt = Checkpoint(self.weights + '.json', auto_update=self.auto_update_checkpoints) restore_checkpoint_params = ckpt.checkpoint restore_model_params = restore_checkpoint_params.model # checks if checkpoint_params.model.line_height != network_params.features: raise Exception("The model to restore has a line height of {} but a line height of {} is requested".format( network_params.features, checkpoint_params.model.line_height )) # create codec of the same type restore_codec = codec.__class__(restore_model_params.codec.charset) # the codec changes as tuple (deletions/insertions), and the new codec is the changed old one codec_changes = restore_codec.align(codec, shrink=not self.keep_loaded_codec) codec = restore_codec print("Codec changes: {} deletions, {} appends".format(len(codec_changes[0]), len(codec_changes[1]))) # The actual weight/bias matrix will be changed after loading the old weights if all([c == 0 for c in codec_changes]): codec_changes = None # No codec changes else: codec_changes = None # store the new codec network_params.classes = len(codec) checkpoint_params.model.codec.charset[:] = codec.charset print("CODEC: {}".format(codec.charset)) backend = create_backend_from_checkpoint( checkpoint_params=checkpoint_params, processes=checkpoint_params.processes, ) train_net = backend.create_net(codec, graph_type="train", checkpoint_to_load=Checkpoint(self.weights) if self.weights else None, batch_size=checkpoint_params.batch_size, codec_changes=codec_changes) if checkpoint_params.current_stage == 0: self._run_train(train_net, train_start_time, progress_bar, self.dataset, self.validation_dataset, training_callback) if checkpoint_params.data_aug_retrain_on_original and self.data_augmenter and self.n_augmentations != 0: print("Starting training on original data only") if checkpoint_params.current_stage == 0: checkpoint_params.current_stage = 1 checkpoint_params.iter = 0 checkpoint_params.early_stopping_best_at_iter = 0 checkpoint_params.early_stopping_best_cur_nbest = 0 checkpoint_params.early_stopping_best_accuracy = 0 self.dataset.generate_only_non_augmented = True # this is the important line! self._run_train(train_net, train_start_time, progress_bar, self.dataset, self.validation_dataset, training_callback)