Esempio n. 1
0
    def final_test(self,
                   base_path: Path,
                   eval_mini_batch_size: int,
                   num_workers: int = 8):

        log_line(log)
        log.info("Testing using best model ...")

        self.model.eval()

        if (base_path / "best-model.pt").exists():
            self.model = self.model.load(base_path / "best-model.pt")

        test_results, test_loss = self.model.evaluate(
            DataLoader(
                self.corpus.test,
                batch_size=eval_mini_batch_size,
                num_workers=num_workers,
            ),
            out_path=base_path / "test.tsv",
            embedding_storage_mode="none",
        )

        test_results: Result = test_results
        log.info(test_results.log_line)
        log.info(test_results.detailed_results)
        log_line(log)

        # if we are training over multiple datasets, do evaluation for each
        if type(self.corpus) is MultiCorpus:
            for subcorpus in self.corpus.corpora:
                log_line(log)
                self.model.evaluate(
                    DataLoader(
                        subcorpus.test,
                        batch_size=eval_mini_batch_size,
                        num_workers=num_workers,
                    ),
                    out_path=base_path / f"{subcorpus.name}-test.tsv",
                    embedding_storage_mode="none",
                )

        # get and return the final test score of best model
        final_score = test_results.main_score

        return final_score
Esempio n. 2
0
def main(data_folder: str, model_folder: str, dev_size: float) -> None:
    nlp = spacy.blank('fr')
    nlp.tokenizer = get_tokenizer(nlp)

    corpus: Corpus = prepare_flair_train_test_corpus(spacy_model=nlp,
                                                     data_folder=data_folder,
                                                     dev_size=dev_size)

    tagger: SequenceTagger = SequenceTagger.load(
        model=os.path.join(model_folder, 'best-model.pt'))

    test_results, _ = tagger.evaluate(data_loader=DataLoader(corpus.test,
                                                             batch_size=32,
                                                             num_workers=10),
                                      embeddings_storage_mode="cpu")
    print(test_results.detailed_results)

    sentences_original = (corpus.train.sentences + corpus.test.sentences)
    sentences_predict = copy.deepcopy(sentences_original)
    # clean tokens in case there is a bug
    for s in sentences_predict:
        for t in s:
            t.tags = {}

    _ = tagger.predict(sentences=sentences_predict,
                       mini_batch_size=32,
                       embedding_storage_mode="cpu",
                       verbose=True)

    entities_to_keep = [
        "PERS", "ADDRESS", "ORGANIZATION", "JUDGE_CLERK", "LAWYER"
    ]
    for index, (sentence_original, sentence_predict) \
            in enumerate(zip(sentences_original, sentences_predict)):  # type: int, (Sentence, Sentence)
        expected_entities_text = {
            f"{s.text} {s.tag}"
            for s in sentence_original.get_spans('ner')
            if s.tag in entities_to_keep
        }
        predicted_entities_text = {
            f"{s.text} {s.tag}"
            for s in sentence_predict.get_spans('ner')
            if s.tag in entities_to_keep if s.score > 0.8
        }

        diff_expected = expected_entities_text.difference(
            predicted_entities_text)
        diff_predicted = predicted_entities_text.difference(
            expected_entities_text)

        if (len(diff_predicted) > 0):  # (len(diff_expected) > 0) or
            print("------------")
            print(f"source {index}: [{sentence_original.to_plain_string()}]")
            print(f"expected missing: [{diff_expected}]")
            print(f"predicted missing: [{diff_predicted}]")
            print(
                f"common: [{set(predicted_entities_text).intersection(set(expected_entities_text))}]"
            )
def main(data_folder: str, model_folder: str, dev_size: float,
         entities_to_remove: List[str]) -> None:
    nlp = spacy.blank('fr')
    nlp.tokenizer = get_tokenizer(nlp)

    corpus: Corpus = prepare_flair_train_test_corpus(spacy_model=nlp,
                                                     data_folder=data_folder,
                                                     dev_size=dev_size)
    # flair.device = torch.device('cpu')  # (4mn 28)
    tagger: SequenceTagger = SequenceTagger.load(
        model=os.path.join(model_folder, 'best-model.pt'))
    test_results, _ = tagger.evaluate(
        data_loader=DataLoader(corpus.test, batch_size=32))
    print(test_results.detailed_results)

    sentences_original = (corpus.train.sentences + corpus.test.sentences)
    sentences_predict = copy.deepcopy(sentences_original)
    # clean tokens in case there is a bug
    for s in sentences_predict:
        for t in s:
            t.tags = {}

    _ = tagger.predict(sentences=sentences_predict,
                       mini_batch_size=32,
                       embedding_storage_mode="none",
                       verbose=True)

    for index, (sentence_original, sentence_predict) \
            in enumerate(zip(sentences_original, sentences_predict)):  # type: int, (Sentence, Sentence)
        expected_entities_text = {
            f"{s.text} {s.tag}"
            for s in sentence_original.get_spans('ner')
            if s.tag not in entities_to_remove
        }
        predicted_entities_text = {
            f"{s.text} {s.tag}"
            for s in sentence_predict.get_spans('ner')
            if s.tag not in entities_to_remove
        }

        diff_expected = expected_entities_text.difference(
            predicted_entities_text)
        diff_predicted = predicted_entities_text.difference(
            expected_entities_text)

        if len(diff_predicted) > 0:  # (len(diff_expected) > 0) or
            print("------------")
            print(f"source {index}: [{sentence_original.to_plain_string()}]")
            print(f"expected missing: [{diff_expected}]")
            print(f"predicted missing: [{diff_predicted}]")
            print(
                f"common: [{set(predicted_entities_text).intersection(set(expected_entities_text))}]"
            )
Esempio n. 4
0
 def _filter_empty_sentences(dataset) -> Dataset:
     empty_sentence_indices = []
     non_empty_sentence_indices = []
     index = 0
     from flair.datasets import DataLoader
     for batch in DataLoader(dataset):
         for sentence in batch:
             if (len(sentence) == 0):
                 empty_sentence_indices.append(index)
             else:
                 non_empty_sentence_indices.append(index)
             index += 1
     subset = Subset(dataset, non_empty_sentence_indices)
     return subset
    def predict(
        self,
        sentences: Union[Sentence, List[Sentence]],
        mini_batch_size: int = 32,
        verbose: bool = False,
        label_name: Optional[str] = None,
        embedding_storage_mode="none",
    ) -> List[Sentence]:

        if label_name is None:
            label_name = self.label_name if self.label_name is not None else "label"

        with torch.no_grad():
            if not isinstance(sentences, list):
                sentences = [sentences]

            if not sentences:
                return sentences

            reordered_sentences = sorted(sentences,
                                         key=lambda s: len(s),
                                         reverse=True)

            if len(reordered_sentences) == 0:
                return sentences

            dataloader = DataLoader(
                dataset=FlairDatapointDataset(reordered_sentences),
                batch_size=mini_batch_size,
            )
            # progress bar for verbosity
            if verbose:
                progress_bar = tqdm(dataloader)
                progress_bar.set_description("Batch inference")
                dataloader = progress_bar

            for batch in dataloader:
                # stop if all sentences are empty
                if not batch:
                    continue
                scores = self.forward(batch)

                for (sentence, score) in zip(batch, scores.tolist()):
                    sentence.set_label(label_name, value=str(score[0]))

                # clearing token embeddings to save memory
                store_embeddings(batch, storage_mode=embedding_storage_mode)

            return sentences
Esempio n. 6
0
 def make_label_dictionary(self) -> Dictionary:
     '\n        Creates a dictionary of all labels assigned to the sentences in the corpus.\n        :return: dictionary of labels\n        '
     label_dictionary = Dictionary(add_unk=False)
     label_dictionary.multi_label = False
     from flair.datasets import DataLoader
     loader = DataLoader(self.train, batch_size=1)
     log.info('Computing label dictionary. Progress:')
     for batch in Tqdm.tqdm(iter(loader)):
         for sentence in batch:
             for label in sentence.labels:
                 label_dictionary.add_item(label.value)
             if (not label_dictionary.multi_label):
                 if (len(sentence.labels) > 1):
                     label_dictionary.multi_label = True
     log.info(label_dictionary.idx2item)
     return label_dictionary
    def predict(
        self,
        sentences: Union[List[Sentence], Sentence],
        mini_batch_size: int = 32,
        num_workers: int = 8,
        print_tree: bool = False,
        embedding_storage_mode="none",
    ) -> None:
        """
        Predict arcs and tags for Dependency Parser task
        :param sentences: a Sentence or a List of Sentence
        :param mini_batch_size: mini batch size to use
        :param print_tree: set to True to print dependency parser of sentence as tree shape
        :param embedding_storage_mode: default is 'none' which is always best. Only set to 'cpu' or 'gpu' if
        you wish to not only predict, but also keep the generated embeddings in CPU or GPU memory respectively.
        'gpu' to store embeddings in GPU memory.
        """
        if not isinstance(sentences, list):
            sentences = [sentences]
        sentence_dataset = FlairDatapointDataset(sentences)
        data_loader = DataLoader(sentence_dataset,
                                 batch_size=mini_batch_size,
                                 num_workers=num_workers)

        for batch in data_loader:
            with torch.no_grad():
                score_arc, score_rel = self.forward(batch)
                arc_prediction, relation_prediction = self._obtain_labels_(
                    score_arc, score_rel)

            for sentnce_index, (sentence, sent_tags, sent_arcs) in enumerate(
                    zip(batch, relation_prediction, arc_prediction)):

                for token_index, (token, tag, head_id) in enumerate(
                        zip(sentence.tokens, sent_tags, sent_arcs)):
                    token.add_tag(self.tag_type, tag,
                                  score_rel[sentnce_index][token_index])

                    token.head_id = int(head_id)

                if print_tree:
                    tree_printer(sentence, self.tag_type)
                    print("-" * 50)
            store_embeddings(batch, storage_mode=embedding_storage_mode)
Esempio n. 8
0
    def _filter_empty_sentences(dataset) -> Dataset:

        # find out empty sentence indices
        empty_sentence_indices = []
        non_empty_sentence_indices = []
        index = 0

        from flair.datasets import DataLoader

        for batch in DataLoader(dataset):
            for sentence in batch:
                if len(sentence) == 0:
                    empty_sentence_indices.append(index)
                else:
                    non_empty_sentence_indices.append(index)
                index += 1

        # create subset of non-empty sentence indices
        subset = Subset(dataset, non_empty_sentence_indices)

        return subset
Esempio n. 9
0
 def format_data(self, dataset_type):
     if dataset_type == "train":  # Flair expects URI paths to data when training.
         data_paths = data_archives.get_dataset_path(self.args.lang,
                                                     self.args.treebank,
                                                     None,
                                                     simplified=False)
         path_only = "/".join(data_paths[0].split("/")[:-1])
         train_file = data_paths[0].split("/")[-1]
         test_file = data_paths[1].split("/")[-1]
         dev_file = data_paths[2].split("/")[-1]
         return (path_only, train_file, test_file, dev_file)
     else:  # Use Flairs built-in DataLoader when loading test data.
         sentences = []
         for batch in DataLoader(self.corpus.test, 32, 8):
             for sentence in batch:
                 curr_sentence = []
                 for token in sentence:
                     true_tag = token.get_tag("upos").value
                     curr_sentence.append((token.text, true_tag))
                 sentences.append(curr_sentence)
         return sentences
Esempio n. 10
0
    def _filter_long_sentences(dataset, max_charlength: int) -> Dataset:

        # find out empty sentence indices
        empty_sentence_indices = []
        non_empty_sentence_indices = []
        index = 0

        from flair.datasets import DataLoader

        for batch in DataLoader(dataset):
            for sentence in batch:
                if len(sentence.to_plain_string()) > max_charlength:
                    empty_sentence_indices.append(index)
                else:
                    non_empty_sentence_indices.append(index)
                index += 1

        # create subset of non-empty sentence indices
        subset = Subset(dataset, non_empty_sentence_indices)

        return subset
Esempio n. 11
0
    def prepare_epoch(self):
        log_line(log)

        for group in self.optimizer.param_groups:
            self.learning_rate = group["lr"]

        # reload last best model if annealing with restarts is enabled
        if (self.learning_rate != self.previous_learning_rate
                and self.anneal_with_restarts
                and (base_path / "best-model.pt").exists()):
            log.info("resetting to best model")
            self.model.load(base_path / "best-model.pt")

        self.previous_learning_rate = self.learning_rate

        # stop training if learning rate becomes too small
        if self.learning_rate < self.min_learning_rate:
            log_line(log)
            log.info("learning rate too small - quitting training!")
            log_line(log)
            return

        self.batch_loader = DataLoader(
            self.train_data,
            batch_size=self.mini_batch_size,
            shuffle=self.shuffle,
            num_workers=self.num_workers,
        )

        self.model.train()
        self.model.set_output(self.model_mode)

        self.train_loss: float = 0
        self.seen_batches = 0
        self.total_number_of_batches = len(self.batch_loader)

        self.modulo = max(1, int(self.total_number_of_batches / 10))
        self.batch_time = 0
Esempio n. 12
0
    def make_label_dictionary(self, label_type: str = None) -> Dictionary:
        """
        Creates a dictionary of all labels assigned to the sentences in the corpus.
        :return: dictionary of labels
        """
        label_dictionary: Dictionary = Dictionary(add_unk=False)
        label_dictionary.multi_label = False

        from flair.datasets import DataLoader

        data = ConcatDataset([self.train, self.test])
        loader = DataLoader(data, batch_size=1)

        log.info("Computing label dictionary. Progress:")
        for batch in Tqdm.tqdm(iter(loader)):

            for sentence in batch:

                # check if sentence itself has labels
                labels = sentence.get_labels(
                    label_type) if label_type is not None else sentence.labels

                for label in labels:
                    label_dictionary.add_item(label.value)

                # check for labels of words
                if isinstance(sentence, Sentence):
                    for token in sentence.tokens:
                        for label in token.get_labels(label_type):
                            label_dictionary.add_item(label.value)

                if not label_dictionary.multi_label:
                    if len(labels) > 1:
                        label_dictionary.multi_label = True

        log.info(label_dictionary.idx2item)

        return label_dictionary
Esempio n. 13
0
    def evaluate(
        self,
        data_points: Union[List[DataPoint], Dataset],
        gold_label_type: str,
        out_path: Union[str, Path] = None,
        embedding_storage_mode: str = "none",
        mini_batch_size: int = 32,
        num_workers: int = 8,
        main_evaluation_metric: Tuple[str, str] = ("micro avg", "f1-score"),
        exclude_labels: List[str] = [],
        gold_label_dictionary: Optional[Dictionary] = None,
    ) -> Result:
        import numpy as np
        import sklearn

        # read Dataset into data loader (if list of sentences passed, make Dataset first)
        if not isinstance(data_points, Dataset):
            data_points = SentenceDataset(data_points)
        data_loader = DataLoader(data_points,
                                 batch_size=mini_batch_size,
                                 num_workers=num_workers)

        with torch.no_grad():

            # loss calculation
            eval_loss = 0
            average_over = 0

            # variables for printing
            lines: List[str] = []

            # variables for computing scores
            all_spans: List[str] = []
            all_true_values = {}
            all_predicted_values = {}

            sentence_id = 0
            for batch in data_loader:

                # remove any previously predicted labels
                for datapoint in batch:
                    datapoint.remove_labels('predicted')

                # predict for batch
                loss_and_count = self.predict(
                    batch,
                    embedding_storage_mode=embedding_storage_mode,
                    mini_batch_size=mini_batch_size,
                    label_name='predicted',
                    return_loss=True)

                if isinstance(loss_and_count, Tuple):
                    average_over += loss_and_count[1]
                    eval_loss += loss_and_count[0]
                else:
                    eval_loss += loss_and_count

                # get the gold labels
                for datapoint in batch:

                    for gold_label in datapoint.get_labels(gold_label_type):
                        representation = str(
                            sentence_id) + ': ' + gold_label.identifier

                        value = gold_label.value
                        if gold_label_dictionary and gold_label_dictionary.get_idx_for_item(
                                value) == 0:
                            value = '<unk>'

                        if representation not in all_true_values:
                            all_true_values[representation] = [value]
                        else:
                            all_true_values[representation].append(value)

                        if representation not in all_spans:
                            all_spans.append(representation)

                    for predicted_span in datapoint.get_labels("predicted"):
                        representation = str(
                            sentence_id) + ': ' + predicted_span.identifier

                        # add to all_predicted_values
                        if representation not in all_predicted_values:
                            all_predicted_values[representation] = [
                                predicted_span.value
                            ]
                        else:
                            all_predicted_values[representation].append(
                                predicted_span.value)

                        if representation not in all_spans:
                            all_spans.append(representation)

                    sentence_id += 1

                store_embeddings(batch, embedding_storage_mode)

                # make printout lines
                if out_path:
                    lines.extend(
                        self._print_predictions(batch, gold_label_type))

            # write all_predicted_values to out_file if set
            if out_path:
                with open(Path(out_path), "w", encoding="utf-8") as outfile:
                    outfile.write("".join(lines))

            # make the evaluation dictionary
            evaluation_label_dictionary = Dictionary(add_unk=False)
            evaluation_label_dictionary.add_item("O")
            for true_values in all_true_values.values():
                for label in true_values:
                    evaluation_label_dictionary.add_item(label)
            for predicted_values in all_predicted_values.values():
                for label in predicted_values:
                    evaluation_label_dictionary.add_item(label)

            # finally, compute numbers
            y_true = []
            y_pred = []

            for span in all_spans:

                true_values = all_true_values[
                    span] if span in all_true_values else ['O']
                predicted_values = all_predicted_values[
                    span] if span in all_predicted_values else ['O']

                y_true_instance = np.zeros(len(evaluation_label_dictionary),
                                           dtype=int)
                for true_value in true_values:
                    y_true_instance[evaluation_label_dictionary.
                                    get_idx_for_item(true_value)] = 1
                y_true.append(y_true_instance.tolist())

                y_pred_instance = np.zeros(len(evaluation_label_dictionary),
                                           dtype=int)
                for predicted_value in predicted_values:
                    y_pred_instance[evaluation_label_dictionary.
                                    get_idx_for_item(predicted_value)] = 1
                y_pred.append(y_pred_instance.tolist())

        # now, calculate evaluation numbers
        target_names = []
        labels = []

        counter = Counter()
        counter.update(
            list(itertools.chain.from_iterable(all_true_values.values())))
        counter.update(
            list(itertools.chain.from_iterable(all_predicted_values.values())))

        for label_name, count in counter.most_common():
            if label_name == 'O': continue
            if label_name in exclude_labels: continue
            target_names.append(label_name)
            labels.append(
                evaluation_label_dictionary.get_idx_for_item(label_name))

        # there is at least one gold label or one prediction (default)
        if len(all_true_values) + len(all_predicted_values) > 1:
            classification_report = sklearn.metrics.classification_report(
                y_true,
                y_pred,
                digits=4,
                target_names=target_names,
                zero_division=0,
                labels=labels,
            )

            classification_report_dict = sklearn.metrics.classification_report(
                y_true,
                y_pred,
                target_names=target_names,
                zero_division=0,
                output_dict=True,
                labels=labels,
            )

            accuracy_score = round(
                sklearn.metrics.accuracy_score(y_true, y_pred), 4)

            precision_score = round(
                classification_report_dict["micro avg"]["precision"], 4)
            recall_score = round(
                classification_report_dict["micro avg"]["recall"], 4)
            micro_f_score = round(
                classification_report_dict["micro avg"]["f1-score"], 4)
            macro_f_score = round(
                classification_report_dict["macro avg"]["f1-score"], 4)

            main_score = classification_report_dict[main_evaluation_metric[0]][
                main_evaluation_metric[1]]

        else:
            # issue error and default all evaluation numbers to 0.
            log.error(
                "ACHTUNG! No gold labels and no all_predicted_values found! Could be an error in your corpus or how you "
                "initialize the trainer!")
            accuracy_score = precision_score = recall_score = micro_f_score = macro_f_score = main_score = 0.
            classification_report = ""
            classification_report_dict = {}

        detailed_result = ("\nResults:"
                           f"\n- F-score (micro) {micro_f_score}"
                           f"\n- F-score (macro) {macro_f_score}"
                           f"\n- Accuracy {accuracy_score}"
                           "\n\nBy class:\n" + classification_report)

        # line for log file
        log_header = "PRECISION\tRECALL\tF1\tACCURACY"
        log_line = f"{precision_score}\t" f"{recall_score}\t" f"{micro_f_score}\t" f"{accuracy_score}"

        if average_over > 0:
            eval_loss /= average_over

        result = Result(main_score=main_score,
                        log_line=log_line,
                        log_header=log_header,
                        detailed_results=detailed_result,
                        classification_report=classification_report_dict,
                        loss=eval_loss)

        return result
Esempio n. 14
0
 def train(self,
           base_path: Union[(Path, str)],
           learning_rate: float = 0.1,
           mini_batch_size: int = 32,
           eval_mini_batch_size: int = None,
           max_epochs: int = 100,
           anneal_factor: float = 0.5,
           patience: int = 3,
           min_learning_rate: float = 0.0001,
           train_with_dev: bool = False,
           monitor_train: bool = False,
           monitor_test: bool = False,
           embeddings_storage_mode: str = 'cpu',
           checkpoint: bool = False,
           save_final_model: bool = True,
           anneal_with_restarts: bool = False,
           shuffle: bool = True,
           param_selection_mode: bool = False,
           num_workers: int = 6,
           sampler=None,
           use_amp: bool = False,
           amp_opt_level: str = 'O1',
           **kwargs) -> dict:
     "\n        Trains any class that implements the flair.nn.Model interface.\n        :param base_path: Main path to which all output during training is logged and models are saved\n        :param learning_rate: Initial learning rate\n        :param mini_batch_size: Size of mini-batches during training\n        :param eval_mini_batch_size: Size of mini-batches during evaluation\n        :param max_epochs: Maximum number of epochs to train. Terminates training if this number is surpassed.\n        :param anneal_factor: The factor by which the learning rate is annealed\n        :param patience: Patience is the number of epochs with no improvement the Trainer waits\n         until annealing the learning rate\n        :param min_learning_rate: If the learning rate falls below this threshold, training terminates\n        :param train_with_dev: If True, training is performed using both train+dev data\n        :param monitor_train: If True, training data is evaluated at end of each epoch\n        :param monitor_test: If True, test data is evaluated at end of each epoch\n        :param embeddings_storage_mode: One of 'none' (all embeddings are deleted and freshly recomputed),\n        'cpu' (embeddings are stored on CPU) or 'gpu' (embeddings are stored on GPU)\n        :param checkpoint: If True, a full checkpoint is saved at end of each epoch\n        :param save_final_model: If True, final model is saved\n        :param anneal_with_restarts: If True, the last best model is restored when annealing the learning rate\n        :param shuffle: If True, data is shuffled during training\n        :param param_selection_mode: If True, testing is performed against dev data. Use this mode when doing\n        parameter selection.\n        :param num_workers: Number of workers in your data loader.\n        :param sampler: You can pass a data sampler here for special sampling of data.\n        :param kwargs: Other arguments for the Optimizer\n        :return:\n        "
     if self.use_tensorboard:
         try:
             from torch.utils.tensorboard import SummaryWriter
             writer = SummaryWriter()
         except:
             log_line(log)
             log.warning(
                 'ATTENTION! PyTorch >= 1.1.0 and pillow are required for TensorBoard support!'
             )
             log_line(log)
             self.use_tensorboard = False
             pass
     if use_amp:
         if (sys.version_info < (3, 0)):
             raise RuntimeError(
                 'Apex currently only supports Python 3. Aborting.')
         if (amp is None):
             raise RuntimeError(
                 'Failed to import apex. Please install apex from https://www.github.com/nvidia/apex to enable mixed-precision training.'
             )
     if (eval_mini_batch_size is None):
         eval_mini_batch_size = mini_batch_size
     if (type(base_path) is str):
         base_path = Path(base_path)
     log_handler = add_file_handler(log, (base_path / 'training.log'))
     log_line(log)
     log.info(''.join(['Model: "', '{}'.format(self.model), '"']))
     log_line(log)
     log.info(''.join(['Corpus: "', '{}'.format(self.corpus), '"']))
     log_line(log)
     log.info('Parameters:')
     log.info(''.join(
         [' - learning_rate: "', '{}'.format(learning_rate), '"']))
     log.info(''.join(
         [' - mini_batch_size: "', '{}'.format(mini_batch_size), '"']))
     log.info(''.join([' - patience: "', '{}'.format(patience), '"']))
     log.info(''.join(
         [' - anneal_factor: "', '{}'.format(anneal_factor), '"']))
     log.info(''.join([' - max_epochs: "', '{}'.format(max_epochs), '"']))
     log.info(''.join([' - shuffle: "', '{}'.format(shuffle), '"']))
     log.info(''.join(
         [' - train_with_dev: "', '{}'.format(train_with_dev), '"']))
     log_line(log)
     log.info(''.join(
         ['Model training base path: "', '{}'.format(base_path), '"']))
     log_line(log)
     log.info(''.join(['Device: ', '{}'.format(flair.device)]))
     log_line(log)
     log.info(''.join([
         'Embeddings storage mode: ', '{}'.format(embeddings_storage_mode)
     ]))
     log_train = (True if monitor_train else False)
     log_test = (True if ((not param_selection_mode) and self.corpus.test
                          and monitor_test) else False)
     log_dev = (True if (not train_with_dev) else False)
     loss_txt = init_output_file(base_path, 'loss.tsv')
     weight_extractor = WeightExtractor(base_path)
     optimizer = self.optimizer(self.model.parameters(),
                                lr=learning_rate,
                                **kwargs)
     if (self.optimizer_state is not None):
         optimizer.load_state_dict(self.optimizer_state)
     if use_amp:
         (self.model, optimizer) = amp.initialize(self.model,
                                                  optimizer,
                                                  opt_level=amp_opt_level)
     anneal_mode = ('min' if train_with_dev else 'max')
     scheduler = ReduceLROnPlateau(optimizer,
                                   factor=anneal_factor,
                                   patience=patience,
                                   mode=anneal_mode,
                                   verbose=True)
     if (self.scheduler_state is not None):
         scheduler.load_state_dict(self.scheduler_state)
     train_data = self.corpus.train
     if train_with_dev:
         train_data = ConcatDataset([self.corpus.train, self.corpus.dev])
     if (sampler is not None):
         sampler = sampler(train_data)
         shuffle = False
     dev_score_history = []
     dev_loss_history = []
     train_loss_history = []
     try:
         previous_learning_rate = learning_rate
         for epoch in range((0 + self.epoch), (max_epochs + self.epoch)):
             log_line(log)
             for group in optimizer.param_groups:
                 learning_rate = group['lr']
             if ((learning_rate != previous_learning_rate)
                     and anneal_with_restarts
                     and (base_path / 'best-model.pt').exists()):
                 log.info('resetting to best model')
                 self.model.load((base_path / 'best-model.pt'))
             previous_learning_rate = learning_rate
             if (learning_rate < min_learning_rate):
                 log_line(log)
                 log.info('learning rate too small - quitting training!')
                 log_line(log)
                 break
             batch_loader = DataLoader(train_data,
                                       batch_size=mini_batch_size,
                                       shuffle=shuffle,
                                       num_workers=num_workers,
                                       sampler=sampler)
             self.model.train()
             train_loss = 0
             seen_batches = 0
             total_number_of_batches = len(batch_loader)
             modulo = max(1, int((total_number_of_batches / 10)))
             batch_time = 0
             for (batch_no, batch) in enumerate(batch_loader):
                 start_time = time.time()
                 loss = self.model.forward_loss(batch)
                 optimizer.zero_grad()
                 if use_amp:
                     with amp.scale_loss(loss, optimizer) as scaled_loss:
                         scaled_loss.backward()
                 else:
                     loss.backward()
                 torch.nn.utils.clip_grad_norm_(self.model.parameters(),
                                                5.0)
                 optimizer.step()
                 seen_batches += 1
                 train_loss += loss.item()
                 store_embeddings(batch, embeddings_storage_mode)
                 batch_time += (time.time() - start_time)
                 if ((batch_no % modulo) == 0):
                     log.info(''.join([
                         'epoch ', '{}'.format((epoch + 1)), ' - iter ',
                         '{}'.format(batch_no), '/',
                         '{}'.format(total_number_of_batches), ' - loss ',
                         '{:.8f}'.format((train_loss / seen_batches)),
                         ' - samples/sec: ', '{:.2f}'.format(
                             ((mini_batch_size * modulo) / batch_time))
                     ]))
                     batch_time = 0
                     iteration = ((epoch * total_number_of_batches) +
                                  batch_no)
                     if (not param_selection_mode):
                         weight_extractor.extract_weights(
                             self.model.state_dict(), iteration)
             train_loss /= seen_batches
             self.model.eval()
             log_line(log)
             log.info(''.join([
                 'EPOCH ', '{}'.format((epoch + 1)), ' done: loss ',
                 '{:.4f}'.format(train_loss), ' - lr ',
                 '{:.4f}'.format(learning_rate)
             ]))
             if self.use_tensorboard:
                 writer.add_scalar('train_loss', train_loss, (epoch + 1))
             current_score = train_loss
             result_line = ''
             if log_train:
                 (train_eval_result, train_loss) = self.model.evaluate(
                     DataLoader(self.corpus.train,
                                batch_size=eval_mini_batch_size,
                                num_workers=num_workers),
                     embeddings_storage_mode=embeddings_storage_mode)
                 result_line += ''.join(
                     ['\t', '{}'.format(train_eval_result.log_line)])
                 store_embeddings(self.corpus.train,
                                  embeddings_storage_mode)
             if log_dev:
                 (dev_eval_result, dev_loss) = self.model.evaluate(
                     DataLoader(self.corpus.dev,
                                batch_size=eval_mini_batch_size,
                                num_workers=num_workers),
                     embeddings_storage_mode=embeddings_storage_mode)
                 result_line += ''.join([
                     '\t', '{}'.format(dev_loss), '\t',
                     '{}'.format(dev_eval_result.log_line)
                 ])
                 log.info(''.join([
                     'DEV : loss ', '{}'.format(dev_loss), ' - score ',
                     '{}'.format(dev_eval_result.main_score)
                 ]))
                 dev_score_history.append(dev_eval_result.main_score)
                 dev_loss_history.append(dev_loss)
                 current_score = dev_eval_result.main_score
                 store_embeddings(self.corpus.dev, embeddings_storage_mode)
                 if self.use_tensorboard:
                     writer.add_scalar('dev_loss', dev_loss, (epoch + 1))
                     writer.add_scalar('dev_score',
                                       dev_eval_result.main_score,
                                       (epoch + 1))
             if log_test:
                 (test_eval_result, test_loss) = self.model.evaluate(
                     DataLoader(self.corpus.test,
                                batch_size=eval_mini_batch_size,
                                num_workers=num_workers),
                     (base_path / 'test.tsv'),
                     embeddings_storage_mode=embeddings_storage_mode)
                 result_line += ''.join([
                     '\t', '{}'.format(test_loss), '\t',
                     '{}'.format(test_eval_result.log_line)
                 ])
                 log.info(''.join([
                     'TEST : loss ', '{}'.format(test_loss), ' - score ',
                     '{}'.format(test_eval_result.main_score)
                 ]))
                 store_embeddings(self.corpus.test, embeddings_storage_mode)
                 if self.use_tensorboard:
                     writer.add_scalar('test_loss', test_loss, (epoch + 1))
                     writer.add_scalar('test_score',
                                       test_eval_result.main_score,
                                       (epoch + 1))
             scheduler.step(current_score)
             train_loss_history.append(train_loss)
             try:
                 bad_epochs = scheduler.num_bad_epochs
             except:
                 bad_epochs = 0
             for group in optimizer.param_groups:
                 new_learning_rate = group['lr']
             if (new_learning_rate != previous_learning_rate):
                 bad_epochs = (patience + 1)
             log.info(''.join(
                 ['BAD EPOCHS (no improvement): ',
                  '{}'.format(bad_epochs)]))
             with open(loss_txt, 'a') as f:
                 if (epoch == 0):
                     f.write(
                         'EPOCH\tTIMESTAMP\tBAD_EPOCHS\tLEARNING_RATE\tTRAIN_LOSS'
                     )
                     if log_train:
                         f.write(('\tTRAIN_' + '\tTRAIN_'.join(
                             train_eval_result.log_header.split('\t'))))
                     if log_dev:
                         f.write(('\tDEV_LOSS\tDEV_' + '\tDEV_'.join(
                             dev_eval_result.log_header.split('\t'))))
                     if log_test:
                         f.write(('\tTEST_LOSS\tTEST_' + '\tTEST_'.join(
                             test_eval_result.log_header.split('\t'))))
                 f.write(''.join([
                     '\n', '{}'.format(epoch), '\t',
                     '{:%H:%M:%S}'.format(datetime.datetime.now()), '\t',
                     '{}'.format(bad_epochs), '\t',
                     '{:.4f}'.format(learning_rate), '\t',
                     '{}'.format(train_loss)
                 ]))
                 f.write(result_line)
             if (checkpoint and (not param_selection_mode)):
                 self.model.save_checkpoint((base_path / 'checkpoint.pt'),
                                            optimizer.state_dict(),
                                            scheduler.state_dict(),
                                            (epoch + 1), train_loss)
             if ((not train_with_dev) and (not param_selection_mode)
                     and (current_score == scheduler.best)):
                 self.model.save((base_path / 'best-model.pt'))
         if (save_final_model and (not param_selection_mode)):
             self.model.save((base_path / 'final-model.pt'))
     except KeyboardInterrupt:
         log_line(log)
         log.info('Exiting from training early.')
         if self.use_tensorboard:
             writer.close()
         if (not param_selection_mode):
             log.info('Saving model ...')
             self.model.save((base_path / 'final-model.pt'))
             log.info('Done.')
     if self.corpus.test:
         final_score = self.final_test(base_path, eval_mini_batch_size,
                                       num_workers)
     else:
         final_score = 0
         log.info('Test data not provided setting final score to 0')
     log.removeHandler(log_handler)
     if self.use_tensorboard:
         writer.close()
     return {
         'test_score': final_score,
         'dev_score_history': dev_score_history,
         'train_loss_history': train_loss_history,
         'dev_loss_history': dev_loss_history,
     }
Esempio n. 15
0
 def find_learning_rate(self,
                        base_path: Union[(Path, str)],
                        file_name: str = 'learning_rate.tsv',
                        start_learning_rate: float = 1e-07,
                        end_learning_rate: float = 10,
                        iterations: int = 100,
                        mini_batch_size: int = 32,
                        stop_early: bool = True,
                        smoothing_factor: float = 0.98,
                        **kwargs) -> Path:
     best_loss = None
     moving_avg_loss = 0
     if (type(base_path) is str):
         base_path = Path(base_path)
     learning_rate_tsv = init_output_file(base_path, file_name)
     with open(learning_rate_tsv, 'a') as f:
         f.write('ITERATION\tTIMESTAMP\tLEARNING_RATE\tTRAIN_LOSS\n')
     optimizer = self.optimizer(self.model.parameters(),
                                lr=start_learning_rate,
                                **kwargs)
     train_data = self.corpus.train
     batch_loader = DataLoader(train_data,
                               batch_size=mini_batch_size,
                               shuffle=True)
     scheduler = ExpAnnealLR(optimizer, end_learning_rate, iterations)
     model_state = self.model.state_dict()
     model_device = next(self.model.parameters()).device
     self.model.train()
     for (itr, batch) in enumerate(batch_loader):
         loss = self.model.forward_loss(batch)
         optimizer.zero_grad()
         loss.backward()
         torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
         optimizer.step()
         scheduler.step(1)
         learning_rate = scheduler.get_lr()[0]
         loss_item = loss.item()
         if (itr == 0):
             best_loss = loss_item
         else:
             if (smoothing_factor > 0):
                 moving_avg_loss = ((smoothing_factor * moving_avg_loss) +
                                    ((1 - smoothing_factor) * loss_item))
                 loss_item = (moving_avg_loss /
                              (1 - (smoothing_factor**(itr + 1))))
             if (loss_item < best_loss):
                 best_loss = loss
         if (stop_early and ((loss_item >
                              (4 * best_loss)) or torch.isnan(loss))):
             log_line(log)
             log.info('loss diverged - stopping early!')
             break
         if (itr > iterations):
             break
         with open(str(learning_rate_tsv), 'a') as f:
             f.write(''.join([
                 '{}'.format(itr), '\t',
                 '{:%H:%M:%S}'.format(datetime.datetime.now()), '\t',
                 '{}'.format(learning_rate), '\t', '{}'.format(loss_item),
                 '\n'
             ]))
     self.model.load_state_dict(model_state)
     self.model.to(model_device)
     log_line(log)
     log.info(''.join([
         'learning rate finder finished - plot ',
         '{}'.format(learning_rate_tsv)
     ]))
     log_line(log)
     return Path(learning_rate_tsv)
Esempio n. 16
0
    def train(
        self,
        base_path: Union[Path, str],
        learning_rate: float = 0.1,
        mini_batch_size: int = 32,
        eval_mini_batch_size: int = None,
        max_epochs: int = 100,
        anneal_factor: float = 0.5,
        patience: int = 3,
        min_learning_rate: float = 0.0001,
        train_with_dev: bool = False,
        monitor_train: bool = False,
        monitor_test: bool = False,
        embeddings_storage_mode: str = "cpu",
        checkpoint: bool = False,
        save_final_model: bool = True,
        anneal_with_restarts: bool = False,
        shuffle: bool = True,
        param_selection_mode: bool = False,
        num_workers: int = 6,
        sampler=None,
        **kwargs,
    ) -> dict:
        """
        Trains any class that implements the flair.nn.Model interface.
        :param base_path: Main path to which all output during training is logged and models are saved
        :param learning_rate: Initial learning rate
        :param mini_batch_size: Size of mini-batches during training
        :param eval_mini_batch_size: Size of mini-batches during evaluation
        :param max_epochs: Maximum number of epochs to train. Terminates training if this number is surpassed.
        :param anneal_factor: The factor by which the learning rate is annealed
        :param patience: Patience is the number of epochs with no improvement the Trainer waits
         until annealing the learning rate
        :param min_learning_rate: If the learning rate falls below this threshold, training terminates
        :param train_with_dev: If True, training is performed using both train+dev data
        :param monitor_train: If True, training data is evaluated at end of each epoch
        :param monitor_test: If True, test data is evaluated at end of each epoch
        :param embeddings_storage_mode: One of 'none' (all embeddings are deleted and freshly recomputed),
        'cpu' (embeddings are stored on CPU) or 'gpu' (embeddings are stored on GPU)
        :param checkpoint: If True, a full checkpoint is saved at end of each epoch
        :param save_final_model: If True, final model is saved
        :param anneal_with_restarts: If True, the last best model is restored when annealing the learning rate
        :param shuffle: If True, data is shuffled during training
        :param param_selection_mode: If True, testing is performed against dev data. Use this mode when doing
        parameter selection.
        :param num_workers: Number of workers in your data loader.
        :param sampler: You can pass a data sampler here for special sampling of data.
        :param kwargs: Other arguments for the Optimizer
        :return:
        """

        if self.use_tensorboard:
            try:
                from torch.utils.tensorboard import SummaryWriter

                writer = SummaryWriter()
            except:
                log_line(log)
                log.warning(
                    "ATTENTION! PyTorch >= 1.1.0 and pillow are required for TensorBoard support!"
                )
                log_line(log)
                self.use_tensorboard = False
                pass

        if eval_mini_batch_size is None:
            eval_mini_batch_size = mini_batch_size

        # cast string to Path
        if type(base_path) is str:
            base_path = Path(base_path)

        log_handler = add_file_handler(log, base_path / "training.log")

        log_line(log)
        log.info(f'Model: "{self.model}"')
        log_line(log)
        log.info(f'Corpus: "{self.corpus}"')
        log_line(log)
        log.info("Parameters:")
        log.info(f' - learning_rate: "{learning_rate}"')
        log.info(f' - mini_batch_size: "{mini_batch_size}"')
        log.info(f' - patience: "{patience}"')
        log.info(f' - anneal_factor: "{anneal_factor}"')
        log.info(f' - max_epochs: "{max_epochs}"')
        log.info(f' - shuffle: "{shuffle}"')
        log.info(f' - train_with_dev: "{train_with_dev}"')
        log_line(log)
        log.info(f'Model training base path: "{base_path}"')
        log_line(log)
        log.info(f"Device: {flair.device}")
        log_line(log)
        log.info(f"Embeddings storage mode: {embeddings_storage_mode}")

        # determine what splits (train, dev, test) to evaluate and log
        log_train = True if monitor_train else False
        log_test = (True if (not param_selection_mode and self.corpus.test
                             and monitor_test) else False)
        log_dev = True if not train_with_dev else False

        # prepare loss logging file and set up header
        loss_txt = init_output_file(base_path, "loss.tsv")

        weight_extractor = WeightExtractor(base_path)

        optimizer: torch.optim.Optimizer = self.optimizer(
            self.model.parameters(), lr=learning_rate, **kwargs)
        if self.optimizer_state is not None:
            optimizer.load_state_dict(self.optimizer_state)

        # minimize training loss if training with dev data, else maximize dev score
        anneal_mode = "min" if train_with_dev else "max"

        scheduler: ReduceLROnPlateau = ReduceLROnPlateau(
            optimizer,
            factor=anneal_factor,
            patience=patience,
            mode=anneal_mode,
            verbose=True,
        )

        if self.scheduler_state is not None:
            scheduler.load_state_dict(self.scheduler_state)

        train_data = self.corpus.train

        # if training also uses dev data, include in training set
        if train_with_dev:
            train_data = ConcatDataset([self.corpus.train, self.corpus.dev])

        if sampler is not None:
            sampler = sampler(train_data)
            shuffle = False

        dev_score_history = []
        dev_loss_history = []
        train_loss_history = []

        # At any point you can hit Ctrl + C to break out of training early.
        try:
            previous_learning_rate = learning_rate

            for epoch in range(0 + self.epoch, max_epochs + self.epoch):
                log_line(log)

                # get new learning rate
                for group in optimizer.param_groups:
                    learning_rate = group["lr"]

                # reload last best model if annealing with restarts is enabled
                if (learning_rate != previous_learning_rate
                        and anneal_with_restarts
                        and (base_path / "best-model.pt").exists()):
                    log.info("resetting to best model")
                    self.model.load(base_path / "best-model.pt")

                previous_learning_rate = learning_rate

                # stop training if learning rate becomes too small
                if learning_rate < min_learning_rate:
                    log_line(log)
                    log.info("learning rate too small - quitting training!")
                    log_line(log)
                    break

                batch_loader = DataLoader(
                    train_data,
                    batch_size=mini_batch_size,
                    shuffle=shuffle,
                    num_workers=num_workers,
                    sampler=sampler,
                )

                self.model.train()

                train_loss: float = 0

                seen_batches = 0
                total_number_of_batches = len(batch_loader)

                modulo = max(1, int(total_number_of_batches / 10))

                # process mini-batches
                for batch_no, batch in enumerate(batch_loader):

                    loss = self.model.forward_loss(batch)

                    optimizer.zero_grad()
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(),
                                                   5.0)
                    optimizer.step()

                    seen_batches += 1
                    train_loss += loss.item()

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(batch, embeddings_storage_mode)

                    if batch_no % modulo == 0:
                        log.info(
                            f"epoch {epoch + 1} - iter {batch_no}/{total_number_of_batches} - loss "
                            f"{train_loss / seen_batches:.8f}")
                        iteration = epoch * total_number_of_batches + batch_no
                        if not param_selection_mode:
                            weight_extractor.extract_weights(
                                self.model.state_dict(), iteration)

                train_loss /= seen_batches

                self.model.eval()

                log_line(log)
                log.info(
                    f"EPOCH {epoch + 1} done: loss {train_loss:.4f} - lr {learning_rate:.4f}"
                )

                if self.use_tensorboard:
                    writer.add_scalar("train_loss", train_loss, epoch + 1)

                # anneal against train loss if training with dev, otherwise anneal against dev score
                current_score = train_loss

                # evaluate on train / dev / test split depending on training settings
                result_line: str = ""

                if log_train:
                    train_eval_result, train_loss = self.model.evaluate(
                        DataLoader(
                            self.corpus.train,
                            batch_size=eval_mini_batch_size,
                            num_workers=num_workers,
                        ),
                        embeddings_storage_mode=embeddings_storage_mode,
                    )
                    result_line += f"\t{train_eval_result.log_line}"

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.train,
                                     embeddings_storage_mode)

                if log_dev:
                    dev_eval_result, dev_loss = self.model.evaluate(
                        DataLoader(
                            self.corpus.dev,
                            batch_size=eval_mini_batch_size,
                            num_workers=num_workers,
                        ),
                        embeddings_storage_mode=embeddings_storage_mode,
                    )
                    result_line += f"\t{dev_loss}\t{dev_eval_result.log_line}"
                    log.info(
                        f"DEV : loss {dev_loss} - score {dev_eval_result.main_score}"
                    )
                    # calculate scores using dev data if available
                    # append dev score to score history
                    dev_score_history.append(dev_eval_result.main_score)
                    dev_loss_history.append(dev_loss)

                    current_score = dev_eval_result.main_score

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.dev, embeddings_storage_mode)

                    if self.use_tensorboard:
                        writer.add_scalar("dev_loss", dev_loss, epoch + 1)
                        writer.add_scalar("dev_score",
                                          dev_eval_result.main_score,
                                          epoch + 1)

                if log_test:
                    test_eval_result, test_loss = self.model.evaluate(
                        DataLoader(
                            self.corpus.test,
                            batch_size=eval_mini_batch_size,
                            num_workers=num_workers,
                        ),
                        base_path / "test.tsv",
                        embeddings_storage_mode=embeddings_storage_mode,
                    )
                    result_line += f"\t{test_loss}\t{test_eval_result.log_line}"
                    log.info(
                        f"TEST : loss {test_loss} - score {test_eval_result.main_score}"
                    )

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.test, embeddings_storage_mode)

                    if self.use_tensorboard:
                        writer.add_scalar("test_loss", test_loss, epoch + 1)
                        writer.add_scalar("test_score",
                                          test_eval_result.main_score,
                                          epoch + 1)

                # determine learning rate annealing through scheduler
                scheduler.step(current_score)

                train_loss_history.append(train_loss)

                # determine bad epoch number
                try:
                    bad_epochs = scheduler.num_bad_epochs
                except:
                    bad_epochs = 0
                for group in optimizer.param_groups:
                    new_learning_rate = group["lr"]
                if new_learning_rate != previous_learning_rate:
                    bad_epochs = patience + 1

                # log bad epochs
                log.info(f"BAD EPOCHS (no improvement): {bad_epochs}")

                # output log file
                with open(loss_txt, "a") as f:

                    # make headers on first epoch
                    if epoch == 0:
                        f.write(
                            f"EPOCH\tTIMESTAMP\tBAD_EPOCHS\tLEARNING_RATE\tTRAIN_LOSS"
                        )

                        if log_train:
                            f.write("\tTRAIN_" + "\tTRAIN_".join(
                                train_eval_result.log_header.split("\t")))
                        if log_dev:
                            f.write("\tDEV_LOSS\tDEV_" + "\tDEV_".join(
                                dev_eval_result.log_header.split("\t")))
                        if log_test:
                            f.write("\tTEST_LOSS\tTEST_" + "\tTEST_".join(
                                test_eval_result.log_header.split("\t")))

                    f.write(
                        f"\n{epoch}\t{datetime.datetime.now():%H:%M:%S}\t{bad_epochs}\t{learning_rate:.4f}\t{train_loss}"
                    )
                    f.write(result_line)

                # if checkpoint is enable, save model at each epoch
                if checkpoint and not param_selection_mode:
                    self.model.save_checkpoint(
                        base_path / "checkpoint.pt",
                        optimizer.state_dict(),
                        scheduler.state_dict(),
                        epoch + 1,
                        train_loss,
                    )

                # if we use dev data, remember best model based on dev evaluation score
                if (not train_with_dev and not param_selection_mode
                        and current_score == scheduler.best):
                    self.model.save(base_path / "best-model.pt")

            # if we do not use dev data for model selection, save final model
            if save_final_model and not param_selection_mode:
                self.model.save(base_path / "final-model.pt")

        except KeyboardInterrupt:
            log_line(log)
            log.info("Exiting from training early.")

            if self.use_tensorboard:
                writer.close()

            if not param_selection_mode:
                log.info("Saving model ...")
                self.model.save(base_path / "final-model.pt")
                log.info("Done.")

        # test best model if test data is present
        if self.corpus.test:
            final_score = self.final_test(base_path, eval_mini_batch_size,
                                          num_workers)
        else:
            final_score = 0
            log.info("Test data not provided setting final score to 0")

        log.removeHandler(log_handler)

        if self.use_tensorboard:
            writer.close()

        return {
            "test_score": final_score,
            "dev_score_history": dev_score_history,
            "train_loss_history": train_loss_history,
            "dev_loss_history": dev_loss_history,
        }
Esempio n. 17
0
    def evaluate(
        self,
        sentences: Union[List[DataPoint], Dataset],
        out_path: Union[str, Path] = None,
        embedding_storage_mode: str = "none",
        mini_batch_size: int = 32,
        num_workers: int = 8,
    ) -> (Result, float):

        # read Dataset into data loader (if list of sentences passed, make Dataset first)
        if not isinstance(sentences, Dataset):
            sentences = SentenceDataset(sentences)
        data_loader = DataLoader(sentences,
                                 batch_size=mini_batch_size,
                                 num_workers=num_workers)

        # use scikit-learn to evaluate
        y_true = []
        y_pred = []

        with torch.no_grad():
            eval_loss = 0

            lines: List[str] = []
            batch_count: int = 0
            for batch in data_loader:

                batch_count += 1

                # remove previously predicted labels
                [sentence.remove_labels('predicted') for sentence in batch]

                # get the gold labels
                true_values_for_batch = [
                    sentence.get_labels(self.label_type) for sentence in batch
                ]

                # predict for batch
                loss = self.predict(
                    batch,
                    embedding_storage_mode=embedding_storage_mode,
                    mini_batch_size=mini_batch_size,
                    label_name='predicted',
                    return_loss=True)

                eval_loss += loss

                sentences_for_batch = [
                    sent.to_plain_string() for sent in batch
                ]

                # get the predicted labels
                predictions = [
                    sentence.get_labels('predicted') for sentence in batch
                ]

                for sentence, prediction, true_value in zip(
                        sentences_for_batch,
                        predictions,
                        true_values_for_batch,
                ):
                    eval_line = "{}\t{}\t{}\n".format(sentence, true_value,
                                                      prediction)
                    lines.append(eval_line)

                for predictions_for_sentence, true_values_for_sentence in zip(
                        predictions, true_values_for_batch):

                    true_values_for_sentence = [
                        label.value for label in true_values_for_sentence
                    ]
                    predictions_for_sentence = [
                        label.value for label in predictions_for_sentence
                    ]

                    y_true_instance = np.zeros(len(self.label_dictionary),
                                               dtype=int)
                    for i in range(len(self.label_dictionary)):
                        if self.label_dictionary.get_item_for_index(
                                i) in true_values_for_sentence:
                            y_true_instance[i] = 1
                    y_true.append(y_true_instance.tolist())

                    y_pred_instance = np.zeros(len(self.label_dictionary),
                                               dtype=int)
                    for i in range(len(self.label_dictionary)):
                        if self.label_dictionary.get_item_for_index(
                                i) in predictions_for_sentence:
                            y_pred_instance[i] = 1
                    y_pred.append(y_pred_instance.tolist())

                store_embeddings(batch, embedding_storage_mode)

            # remove predicted labels
            for sentence in sentences:
                sentence.annotation_layers['predicted'] = []

            if out_path is not None:
                with open(out_path, "w", encoding="utf-8") as outfile:
                    outfile.write("".join(lines))

            # make "classification report"
            target_names = []
            for i in range(len(self.label_dictionary)):
                target_names.append(
                    self.label_dictionary.get_item_for_index(i))
            classification_report = metrics.classification_report(
                y_true,
                y_pred,
                digits=4,
                target_names=target_names,
                zero_division=0)

            # get scores
            micro_f_score = round(
                metrics.fbeta_score(y_true,
                                    y_pred,
                                    beta=self.beta,
                                    average='micro',
                                    zero_division=0), 4)
            accuracy_score = round(metrics.accuracy_score(y_true, y_pred), 4)
            macro_f_score = round(
                metrics.fbeta_score(y_true,
                                    y_pred,
                                    beta=self.beta,
                                    average='macro',
                                    zero_division=0), 4)
            precision_score = round(
                metrics.precision_score(y_true,
                                        y_pred,
                                        average='macro',
                                        zero_division=0), 4)
            recall_score = round(
                metrics.recall_score(y_true,
                                     y_pred,
                                     average='macro',
                                     zero_division=0), 4)

            detailed_result = ("\nResults:"
                               f"\n- F-score (micro) {micro_f_score}"
                               f"\n- F-score (macro) {macro_f_score}"
                               f"\n- Accuracy {accuracy_score}"
                               '\n\nBy class:\n' + classification_report)

            # line for log file
            if not self.multi_label:
                log_header = "ACCURACY"
                log_line = f"\t{accuracy_score}"
            else:
                log_header = "PRECISION\tRECALL\tF1\tACCURACY"
                log_line = f"{precision_score}\t" \
                           f"{recall_score}\t" \
                           f"{macro_f_score}\t" \
                           f"{accuracy_score}"

            result = Result(
                main_score=micro_f_score,
                log_line=log_line,
                log_header=log_header,
                detailed_results=detailed_result,
            )

            eval_loss /= batch_count

            return result, eval_loss
Esempio n. 18
0
    def evaluate(
        self,
        sentences: Dataset,
        eval_mini_batch_size: int = 32,
        embeddings_in_memory: bool = True,
        out_path: Path = None,
        num_workers: int = 8,
    ) -> (Result, float):

        with torch.no_grad():
            eval_loss = 0

            batch_no: int = 0

            batch_loader = DataLoader(
                sentences,
                batch_size=eval_mini_batch_size,
                shuffle=False,
                num_workers=num_workers,
            )

            metric = Metric("Evaluation")

            lines: List[str] = []
            for batch in batch_loader:
                batch_no += 1

                with torch.no_grad():
                    features = self.forward(batch)
                    loss = self._calculate_loss(features, batch)
                    tags, _ = self._obtain_labels(features, batch)

                eval_loss += loss

                for (sentence, sent_tags) in zip(batch, tags):
                    for (token, tag) in zip(sentence.tokens, sent_tags):
                        token: Token = token
                        token.add_tag_label("predicted", tag)

                        # append both to file for evaluation
                        eval_line = "{} {} {} {}\n".format(
                            token.text,
                            token.get_tag(self.tag_type).value,
                            tag.value,
                            tag.score,
                        )
                        lines.append(eval_line)
                    lines.append("\n")
                for sentence in batch:
                    # make list of gold tags
                    gold_tags = [
                        (tag.tag, str(tag)) for tag in sentence.get_spans(self.tag_type)
                    ]
                    # make list of predicted tags
                    predicted_tags = [
                        (tag.tag, str(tag)) for tag in sentence.get_spans("predicted")
                    ]

                    # check for true positives, false positives and false negatives
                    for tag, prediction in predicted_tags:
                        if (tag, prediction) in gold_tags:
                            metric.add_tp(tag)
                        else:
                            metric.add_fp(tag)

                    for tag, gold in gold_tags:
                        if (tag, gold) not in predicted_tags:
                            metric.add_fn(tag)
                        else:
                            metric.add_tn(tag)

                clear_embeddings(
                    batch, also_clear_word_embeddings=not embeddings_in_memory
                )

            eval_loss /= batch_no

            if out_path is not None:
                with open(out_path, "w", encoding="utf-8") as outfile:
                    outfile.write("".join(lines))

            detailed_result = (
                f"\nMICRO_AVG: acc {metric.micro_avg_accuracy()} - f1-score {metric.micro_avg_f_score()}"
                f"\nMACRO_AVG: acc {metric.macro_avg_accuracy()} - f1-score {metric.macro_avg_f_score()}"
            )
            for class_name in metric.get_classes():
                detailed_result += (
                    f"\n{class_name:<10} tp: {metric.get_tp(class_name)} - fp: {metric.get_fp(class_name)} - "
                    f"fn: {metric.get_fn(class_name)} - tn: {metric.get_tn(class_name)} - precision: "
                    f"{metric.precision(class_name):.4f} - recall: {metric.recall(class_name):.4f} - "
                    f"accuracy: {metric.accuracy(class_name):.4f} - f1-score: "
                    f"{metric.f_score(class_name):.4f}"
                )

            result = Result(
                main_score=metric.micro_avg_f_score(),
                log_line=f"{metric.precision()}\t{metric.recall()}\t{metric.micro_avg_f_score()}",
                log_header="PRECISION\tRECALL\tF1",
                detailed_results=detailed_result,
            )

            return result, eval_loss
Esempio n. 19
0
    def predict(
        self,
        sentences: Union[List[Sentence], Sentence],
        mini_batch_size: int = 32,
        multi_class_prob: bool = False,
        verbose: bool = False,
        label_name: Optional[str] = None,
        return_loss=False,
        embedding_storage_mode="none",
    ):
        """
        Predicts the class labels for the given sentences. The labels are directly added to the sentences.
        :param sentences: list of sentences
        :param mini_batch_size: mini batch size to use
        :param multi_class_prob : return probability for all class for multiclass
        :param verbose: set to True to display a progress bar
        :param return_loss: set to True to return loss
        :param label_name: set this to change the name of the label type that is predicted
        :param embedding_storage_mode: default is 'none' which is always best. Only set to 'cpu' or 'gpu' if
        you wish to not only predict, but also keep the generated embeddings in CPU or GPU memory respectively.
        'gpu' to store embeddings in GPU memory.
        """
        if label_name == None:
            label_name = self.label_type if self.label_type is not None else 'label'

        with torch.no_grad():
            if not sentences:
                return sentences

            if isinstance(sentences, DataPoint):
                sentences = [sentences]

            # filter empty sentences
            if isinstance(sentences[0], Sentence):
                sentences = [
                    sentence for sentence in sentences if len(sentence) > 0
                ]
            if len(sentences) == 0: return sentences

            # reverse sort all sequences by their length
            rev_order_len_index = sorted(range(len(sentences)),
                                         key=lambda k: len(sentences[k]),
                                         reverse=True)

            reordered_sentences: List[Union[DataPoint, str]] = [
                sentences[index] for index in rev_order_len_index
            ]

            dataloader = DataLoader(
                dataset=SentenceDataset(reordered_sentences),
                batch_size=mini_batch_size)
            # progress bar for verbosity
            if verbose:
                dataloader = tqdm(dataloader)

            overall_loss = 0
            batch_no = 0
            for batch in dataloader:

                batch_no += 1

                if verbose:
                    dataloader.set_description(
                        f"Inferencing on batch {batch_no}")

                # stop if all sentences are empty
                if not batch:
                    continue

                scores = self.forward(batch)

                if return_loss:
                    overall_loss += self._calculate_loss(scores, batch)

                predicted_labels = self._obtain_labels(
                    scores, predict_prob=multi_class_prob)

                for (sentence, labels) in zip(batch, predicted_labels):
                    for label in labels:
                        if self.multi_label or multi_class_prob:
                            sentence.add_label(label_name, label.value,
                                               label.score)
                        else:
                            sentence.set_label(label_name, label.value,
                                               label.score)

                # clearing token embeddings to save memory
                store_embeddings(batch, storage_mode=embedding_storage_mode)

            if return_loss:
                return overall_loss / batch_no
Esempio n. 20
0
    def predict(
            self,
            sentences: Union[List[Sentence], Sentence],
            mini_batch_size=32,
            all_tag_prob: bool = False,
            verbose: bool = False,
            label_name: Optional[str] = None,
            return_loss=False,
            embedding_storage_mode="none",
    ):
        """
        Predict sequence tags for Named Entity Recognition task
        :param sentences: a Sentence or a List of Sentence
        :param mini_batch_size: size of the minibatch, usually bigger is more rapid but consume more memory,
        up to a point when it has no more effect.
        :param all_tag_prob: True to compute the score for each tag on each token,
        otherwise only the score of the best tag is returned
        :param verbose: set to True to display a progress bar
        :param return_loss: set to True to return loss
        :param label_name: set this to change the name of the label type that is predicted
        :param embedding_storage_mode: default is 'none' which is always best. Only set to 'cpu' or 'gpu' if
        you wish to not only predict, but also keep the generated embeddings in CPU or GPU memory respectively.
        'gpu' to store embeddings in GPU memory.
        """
        if label_name == None:
            label_name = self.tag_type

        with torch.no_grad():
            if not sentences:
                return sentences

            if isinstance(sentences, Sentence):
                sentences = [sentences]

            # reverse sort all sequences by their length
            rev_order_len_index = sorted(
                range(len(sentences)), key=lambda k: len(sentences[k]), reverse=True
            )

            reordered_sentences: List[Union[Sentence, str]] = [
                sentences[index] for index in rev_order_len_index
            ]

            dataloader = DataLoader(
                dataset=SentenceDataset(reordered_sentences), batch_size=mini_batch_size
            )

            if self.use_crf:
                transitions = self.transitions.detach().cpu().numpy()
            else:
                transitions = None

            # progress bar for verbosity
            if verbose:
                dataloader = tqdm(dataloader)

            overall_loss = 0
            batch_no = 0
            for batch in dataloader:

                batch_no += 1

                if verbose:
                    dataloader.set_description(f"Inferencing on batch {batch_no}")

                batch = self._filter_empty_sentences(batch)
                # stop if all sentences are empty
                if not batch:
                    continue

                feature = self.forward(batch)

                if return_loss:
                    overall_loss += self._calculate_loss(feature, batch)

                tags, all_tags = self._obtain_labels(
                    feature=feature,
                    batch_sentences=batch,
                    transitions=transitions,
                    get_all_tags=all_tag_prob,
                )

                for (sentence, sent_tags) in zip(batch, tags):
                    for (token, tag) in zip(sentence.tokens, sent_tags):
                        token.add_tag_label(label_name, tag)

                # all_tags will be empty if all_tag_prob is set to False, so the for loop will be avoided
                for (sentence, sent_all_tags) in zip(batch, all_tags):
                    for (token, token_all_tags) in zip(sentence.tokens, sent_all_tags):
                        token.add_tags_proba_dist(label_name, token_all_tags)

                # clearing token embeddings to save memory
                store_embeddings(batch, storage_mode=embedding_storage_mode)

            if return_loss:
                return overall_loss / batch_no
Esempio n. 21
0
    def evaluate(
            self,
            sentences: Union[List[Sentence], Dataset],
            out_path: Union[str, Path] = None,
            embedding_storage_mode: str = "none",
            mini_batch_size: int = 32,
            num_workers: int = 8,
    ) -> (Result, float):

        # read Dataset into data loader (if list of sentences passed, make Dataset first)
        if not isinstance(sentences, Dataset):
            sentences = SentenceDataset(sentences)
        data_loader = DataLoader(sentences, batch_size=mini_batch_size, num_workers=num_workers)

        # if span F1 needs to be used, use separate eval method
        if self._requires_span_F1_evaluation():
            return self._evaluate_with_span_F1(data_loader, embedding_storage_mode, mini_batch_size, out_path)

        # else, use scikit-learn to evaluate
        y_true = []
        y_pred = []
        labels = Dictionary(add_unk=False)

        eval_loss = 0
        batch_no: int = 0

        lines: List[str] = []

        for batch in data_loader:

            # predict for batch
            loss = self.predict(batch,
                                embedding_storage_mode=embedding_storage_mode,
                                mini_batch_size=mini_batch_size,
                                label_name='predicted',
                                return_loss=True)
            eval_loss += loss
            batch_no += 1

            for sentence in batch:

                for token in sentence:
                    # add gold tag
                    gold_tag = token.get_tag(self.tag_type).value
                    y_true.append(labels.add_item(gold_tag))

                    # add predicted tag
                    predicted_tag = token.get_tag('predicted').value
                    y_pred.append(labels.add_item(predicted_tag))

                    # for file output
                    lines.append(f'{token.text} {gold_tag} {predicted_tag}\n')

                lines.append('\n')

        if out_path:
            with open(Path(out_path), "w", encoding="utf-8") as outfile:
                outfile.write("".join(lines))

        eval_loss /= batch_no

        # use sklearn
        from sklearn import metrics

        # make "classification report"
        target_names = []
        for i in range(len(labels)):
            target_names.append(labels.get_item_for_index(i))
        classification_report = metrics.classification_report(y_true, y_pred, digits=4, target_names=target_names,
                                                              zero_division=1)

        # get scores
        macro_f_score = round(metrics.fbeta_score(y_true, y_pred, beta=self.beta, average='micro'), 4)
        micro_f_score = round(metrics.fbeta_score(y_true, y_pred, beta=self.beta, average='macro'), 4)
        accuracy_score = round(metrics.accuracy_score(y_true, y_pred), 4)

        detailed_result = (
                "\nResults:"
                f"\n- F-score (micro) {macro_f_score}"
                f"\n- F-score (macro) {micro_f_score}"
                f"\n- Accuracy {accuracy_score}"
                '\n\nBy class:\n' + classification_report
        )

        # line for log file
        log_header = "ACCURACY"
        log_line = f"\t{accuracy_score}"

        result = Result(
            main_score=macro_f_score,
            log_line=log_line,
            log_header=log_header,
            detailed_results=detailed_result,
        )
        return result, eval_loss
Esempio n. 22
0
    def train(
        self,
        base_path: Union[Path, str],
        learning_rate: float = 0.1,
        mini_batch_size: int = 32,
        mini_batch_chunk_size: int = None,
        max_epochs: int = 100,
        scheduler=AnnealOnPlateau,
        cycle_momentum: bool = False,
        anneal_factor: float = 0.5,
        patience: int = 3,
        initial_extra_patience=0,
        min_learning_rate: float = 0.0001,
        train_with_dev: bool = False,
        train_with_test: bool = False,
        monitor_train: bool = False,
        monitor_test: bool = False,
        embeddings_storage_mode: str = "cpu",
        checkpoint: bool = False,
        save_final_model: bool = True,
        anneal_with_restarts: bool = False,
        anneal_with_prestarts: bool = False,
        batch_growth_annealing: bool = False,
        shuffle: bool = True,
        param_selection_mode: bool = False,
        write_weights: bool = False,
        num_workers: int = 6,
        sampler=None,
        use_amp: bool = False,
        amp_opt_level: str = "O1",
        eval_on_train_fraction=0.0,
        eval_on_train_shuffle=False,
        save_model_at_each_epoch=False,
        **kwargs,
    ) -> dict:
        """
        Trains any class that implements the flair.nn.Model interface.
        :param base_path: Main path to which all output during training is logged and models are saved
        :param learning_rate: Initial learning rate (or max, if scheduler is OneCycleLR)
        :param mini_batch_size: Size of mini-batches during training
        :param mini_batch_chunk_size: If mini-batches are larger than this number, they get broken down into chunks of this size for processing purposes
        :param max_epochs: Maximum number of epochs to train. Terminates training if this number is surpassed.
        :param scheduler: The learning rate scheduler to use
        :param cycle_momentum: If scheduler is OneCycleLR, whether the scheduler should cycle also the momentum
        :param anneal_factor: The factor by which the learning rate is annealed
        :param patience: Patience is the number of epochs with no improvement the Trainer waits
         until annealing the learning rate
        :param min_learning_rate: If the learning rate falls below this threshold, training terminates
        :param train_with_dev: If True, training is performed using both train+dev data
        :param monitor_train: If True, training data is evaluated at end of each epoch
        :param monitor_test: If True, test data is evaluated at end of each epoch
        :param embeddings_storage_mode: One of 'none' (all embeddings are deleted and freshly recomputed),
        'cpu' (embeddings are stored on CPU) or 'gpu' (embeddings are stored on GPU)
        :param checkpoint: If True, a full checkpoint is saved at end of each epoch
        :param save_final_model: If True, final model is saved
        :param anneal_with_restarts: If True, the last best model is restored when annealing the learning rate
        :param shuffle: If True, data is shuffled during training
        :param param_selection_mode: If True, testing is performed against dev data. Use this mode when doing
        parameter selection.
        :param num_workers: Number of workers in your data loader.
        :param sampler: You can pass a data sampler here for special sampling of data.
        :param eval_on_train_fraction: the fraction of train data to do the evaluation on,
        if 0. the evaluation is not performed on fraction of training data,
        if 'dev' the size is determined from dev set size
        :param eval_on_train_shuffle: if True the train data fraction is determined on the start of training
        and kept fixed during training, otherwise it's sampled at beginning of each epoch
        :param save_model_at_each_epoch: If True, at each epoch the thus far trained model will be saved
        :param kwargs: Other arguments for the Optimizer
        :return:
        """

        if self.use_tensorboard:
            try:
                from torch.utils.tensorboard import SummaryWriter

                writer = SummaryWriter()
            except:
                log_line(log)
                log.warning(
                    "ATTENTION! PyTorch >= 1.1.0 and pillow are required for TensorBoard support!"
                )
                log_line(log)
                self.use_tensorboard = False
                pass

        if use_amp:
            if sys.version_info < (3, 0):
                raise RuntimeError(
                    "Apex currently only supports Python 3. Aborting.")
            if amp is None:
                raise RuntimeError(
                    "Failed to import apex. Please install apex from https://www.github.com/nvidia/apex "
                    "to enable mixed-precision training.")

        if mini_batch_chunk_size is None:
            mini_batch_chunk_size = mini_batch_size
        if learning_rate < min_learning_rate:
            min_learning_rate = learning_rate / 10

        initial_learning_rate = learning_rate

        # cast string to Path
        if type(base_path) is str:
            base_path = Path(base_path)

        log_handler = add_file_handler(log, base_path / "training.log")

        log_line(log)
        log.info(f'Model: "{self.model}"')
        log_line(log)
        log.info(f'Corpus: "{self.corpus}"')
        log_line(log)
        log.info("Parameters:")
        log.info(f' - learning_rate: "{learning_rate}"')
        log.info(f' - mini_batch_size: "{mini_batch_size}"')
        log.info(f' - patience: "{patience}"')
        log.info(f' - anneal_factor: "{anneal_factor}"')
        log.info(f' - max_epochs: "{max_epochs}"')
        log.info(f' - shuffle: "{shuffle}"')
        log.info(f' - train_with_dev: "{train_with_dev}"')
        log.info(f' - batch_growth_annealing: "{batch_growth_annealing}"')
        log_line(log)
        log.info(f'Model training base path: "{base_path}"')
        log_line(log)
        log.info(f"Device: {flair.device}")
        log_line(log)
        log.info(f"Embeddings storage mode: {embeddings_storage_mode}")
        if isinstance(self.model, SequenceTagger
                      ) and self.model.weight_dict and self.model.use_crf:
            log_line(log)
            log.warning(
                f'WARNING: Specified class weights will not take effect when using CRF'
            )

        # determine what splits (train, dev, test) to evaluate and log
        log_train = True if monitor_train else False
        log_test = (True if (not param_selection_mode and self.corpus.test
                             and monitor_test) else False)
        log_dev = False if train_with_dev or not self.corpus.dev else True
        log_train_part = (True if (eval_on_train_fraction == "dev"
                                   or eval_on_train_fraction > 0.0) else False)

        if log_train_part:
            train_part_size = (len(
                self.corpus.dev) if eval_on_train_fraction == "dev" else int(
                    len(self.corpus.train) * eval_on_train_fraction))
            assert train_part_size > 0
            if not eval_on_train_shuffle:
                train_part_indices = list(range(train_part_size))
                train_part = torch.utils.data.dataset.Subset(
                    self.corpus.train, train_part_indices)

        # prepare loss logging file and set up header
        loss_txt = init_output_file(base_path, "loss.tsv")

        weight_extractor = WeightExtractor(base_path)

        optimizer: torch.optim.Optimizer = self.optimizer(
            self.model.parameters(), lr=learning_rate, **kwargs)

        if use_amp:
            self.model, optimizer = amp.initialize(self.model,
                                                   optimizer,
                                                   opt_level=amp_opt_level)

        # minimize training loss if training with dev data, else maximize dev score
        anneal_mode = "min" if train_with_dev else "max"

        if scheduler == OneCycleLR:
            dataset_size = len(self.corpus.train)
            if train_with_dev:
                dataset_size += len(self.corpus.dev)
            lr_scheduler = OneCycleLR(
                optimizer,
                max_lr=learning_rate,
                steps_per_epoch=dataset_size // mini_batch_size + 1,
                epochs=max_epochs - self.
                epoch,  # if we load a checkpoint, we have already trained for self.epoch
                pct_start=0.0,
                cycle_momentum=cycle_momentum)
        else:
            lr_scheduler = scheduler(
                optimizer,
                factor=anneal_factor,
                patience=patience,
                initial_extra_patience=initial_extra_patience,
                mode=anneal_mode,
                verbose=True,
            )

        if (isinstance(lr_scheduler, OneCycleLR) and batch_growth_annealing):
            raise ValueError(
                "Batch growth with OneCycle policy is not implemented.")

        train_data = self.corpus.train

        # if training also uses dev/train data, include in training set
        if train_with_dev or train_with_test:

            parts = [self.corpus.train]
            if train_with_dev: parts.append(self.corpus.dev)
            if train_with_test: parts.append(self.corpus.test)

            train_data = ConcatDataset(parts)

        # initialize sampler if provided
        if sampler is not None:
            # init with default values if only class is provided
            if inspect.isclass(sampler):
                sampler = sampler()
            # set dataset to sample from
            sampler.set_dataset(train_data)
            shuffle = False

        dev_score_history = []
        dev_loss_history = []
        train_loss_history = []

        micro_batch_size = mini_batch_chunk_size

        # At any point you can hit Ctrl + C to break out of training early.
        try:
            previous_learning_rate = learning_rate
            momentum = 0
            for group in optimizer.param_groups:
                if "momentum" in group:
                    momentum = group["momentum"]

            for self.epoch in range(self.epoch + 1, max_epochs + 1):
                log_line(log)

                if anneal_with_prestarts:
                    last_epoch_model_state_dict = copy.deepcopy(
                        self.model.state_dict())

                if eval_on_train_shuffle:
                    train_part_indices = list(range(self.corpus.train))
                    random.shuffle(train_part_indices)
                    train_part_indices = train_part_indices[:train_part_size]
                    train_part = torch.utils.data.dataset.Subset(
                        self.corpus.train, train_part_indices)

                # get new learning rate
                for group in optimizer.param_groups:
                    learning_rate = group["lr"]

                if learning_rate != previous_learning_rate and batch_growth_annealing:
                    mini_batch_size *= 2

                # reload last best model if annealing with restarts is enabled
                if ((anneal_with_restarts or anneal_with_prestarts)
                        and learning_rate != previous_learning_rate
                        and (base_path / "best-model.pt").exists()):
                    if anneal_with_restarts:
                        log.info("resetting to best model")
                        self.model.load_state_dict(
                            self.model.load(base_path /
                                            "best-model.pt").state_dict())
                    if anneal_with_prestarts:
                        log.info("resetting to pre-best model")
                        self.model.load_state_dict(
                            self.model.load(base_path /
                                            "pre-best-model.pt").state_dict())

                previous_learning_rate = learning_rate

                # stop training if learning rate becomes too small
                if (not isinstance(lr_scheduler, OneCycleLR)
                    ) and learning_rate < min_learning_rate:
                    log_line(log)
                    log.info("learning rate too small - quitting training!")
                    log_line(log)
                    break

                batch_loader = DataLoader(
                    train_data,
                    batch_size=mini_batch_size,
                    shuffle=shuffle,
                    num_workers=num_workers,
                    sampler=sampler,
                )

                self.model.train()

                train_loss: float = 0

                seen_batches = 0
                total_number_of_batches = len(batch_loader)

                modulo = max(1, int(total_number_of_batches / 10))

                # process mini-batches
                batch_time = 0
                for batch_no, batch in enumerate(batch_loader):
                    start_time = time.time()

                    # zero the gradients on the model and optimizer
                    self.model.zero_grad()
                    optimizer.zero_grad()

                    # if necessary, make batch_steps
                    batch_steps = [batch]
                    if len(batch) > micro_batch_size:
                        batch_steps = [
                            batch[x:x + micro_batch_size]
                            for x in range(0, len(batch), micro_batch_size)
                        ]

                    # forward and backward for batch
                    for batch_step in batch_steps:

                        # forward pass
                        loss = self.model.forward_loss(batch_step)

                        # Backward
                        if use_amp:
                            with amp.scale_loss(loss,
                                                optimizer) as scaled_loss:
                                scaled_loss.backward()
                        else:
                            loss.backward()

                    # do the optimizer step
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(),
                                                   5.0)
                    optimizer.step()

                    # do the scheduler step if one-cycle
                    if isinstance(lr_scheduler, OneCycleLR):
                        lr_scheduler.step()
                        # get new learning rate
                        for group in optimizer.param_groups:
                            learning_rate = group["lr"]
                            if "momentum" in group:
                                momentum = group["momentum"]

                    seen_batches += 1
                    train_loss += loss.item()

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(batch, embeddings_storage_mode)

                    batch_time += time.time() - start_time
                    if seen_batches % modulo == 0:
                        momentum_info = f' - momentum: {momentum:.4f}' if cycle_momentum else ''
                        log.info(
                            f"epoch {self.epoch} - iter {seen_batches}/{total_number_of_batches} - loss "
                            f"{train_loss / seen_batches:.8f} - samples/sec: {mini_batch_size * modulo / batch_time:.2f}"
                            f" - lr: {learning_rate:.6f}{momentum_info}")
                        batch_time = 0
                        iteration = self.epoch * total_number_of_batches + batch_no
                        if not param_selection_mode and write_weights:
                            weight_extractor.extract_weights(
                                self.model.state_dict(), iteration)

                train_loss /= seen_batches

                self.model.eval()

                log_line(log)
                log.info(
                    f"EPOCH {self.epoch} done: loss {train_loss:.4f} - lr {learning_rate:.7f}"
                )

                if self.use_tensorboard:
                    writer.add_scalar("train_loss", train_loss, self.epoch)

                # anneal against train loss if training with dev, otherwise anneal against dev score
                current_score = train_loss

                # evaluate on train / dev / test split depending on training settings
                result_line: str = ""

                if log_train:
                    train_eval_result, train_loss = self.model.evaluate(
                        self.corpus.train,
                        mini_batch_size=mini_batch_chunk_size,
                        num_workers=num_workers,
                        embedding_storage_mode=embeddings_storage_mode,
                    )
                    result_line += f"\t{train_eval_result.log_line}"

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.train,
                                     embeddings_storage_mode)

                if log_train_part:
                    train_part_eval_result, train_part_loss = self.model.evaluate(
                        train_part,
                        mini_batch_size=mini_batch_chunk_size,
                        num_workers=num_workers,
                        embedding_storage_mode=embeddings_storage_mode,
                    )
                    result_line += (
                        f"\t{train_part_loss}\t{train_part_eval_result.log_line}"
                    )
                    log.info(
                        f"TRAIN_SPLIT : loss {train_part_loss} - score {round(train_part_eval_result.main_score, 4)}"
                    )

                if log_dev:
                    dev_eval_result, dev_loss = self.model.evaluate(
                        self.corpus.dev,
                        mini_batch_size=mini_batch_chunk_size,
                        num_workers=num_workers,
                        out_path=base_path / "dev.tsv",
                        embedding_storage_mode=embeddings_storage_mode,
                    )
                    result_line += f"\t{dev_loss}\t{dev_eval_result.log_line}"
                    log.info(
                        f"DEV : loss {dev_loss} - score {round(dev_eval_result.main_score, 4)}"
                    )
                    # calculate scores using dev data if available
                    # append dev score to score history
                    dev_score_history.append(dev_eval_result.main_score)
                    dev_loss_history.append(dev_loss.item())

                    current_score = dev_eval_result.main_score

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.dev, embeddings_storage_mode)

                    if self.use_tensorboard:
                        writer.add_scalar("dev_loss", dev_loss, self.epoch)
                        writer.add_scalar("dev_score",
                                          dev_eval_result.main_score,
                                          self.epoch)

                if log_test:
                    test_eval_result, test_loss = self.model.evaluate(
                        self.corpus.test,
                        mini_batch_size=mini_batch_chunk_size,
                        num_workers=num_workers,
                        out_path=base_path / "test.tsv",
                        embedding_storage_mode=embeddings_storage_mode,
                    )
                    result_line += f"\t{test_loss}\t{test_eval_result.log_line}"
                    log.info(
                        f"TEST : loss {test_loss} - score {round(test_eval_result.main_score, 4)}"
                    )

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.test, embeddings_storage_mode)

                    if self.use_tensorboard:
                        writer.add_scalar("test_loss", test_loss, self.epoch)
                        writer.add_scalar("test_score",
                                          test_eval_result.main_score,
                                          self.epoch)

                # determine learning rate annealing through scheduler. Use auxiliary metric for AnnealOnPlateau
                if log_dev and isinstance(lr_scheduler, AnnealOnPlateau):
                    lr_scheduler.step(current_score, dev_loss)
                elif not isinstance(lr_scheduler, OneCycleLR):
                    lr_scheduler.step(current_score)

                train_loss_history.append(train_loss)

                # determine bad epoch number
                try:
                    bad_epochs = lr_scheduler.num_bad_epochs
                except:
                    bad_epochs = 0
                for group in optimizer.param_groups:
                    new_learning_rate = group["lr"]
                if new_learning_rate != previous_learning_rate:
                    bad_epochs = patience + 1
                    if previous_learning_rate == initial_learning_rate:
                        bad_epochs += initial_extra_patience

                # log bad epochs
                log.info(f"BAD EPOCHS (no improvement): {bad_epochs}")

                # output log file
                with open(loss_txt, "a") as f:

                    # make headers on first epoch
                    if self.epoch == 1:
                        f.write(
                            f"EPOCH\tTIMESTAMP\tBAD_EPOCHS\tLEARNING_RATE\tTRAIN_LOSS"
                        )

                        if log_train:
                            f.write("\tTRAIN_" + "\tTRAIN_".join(
                                train_eval_result.log_header.split("\t")))
                        if log_train_part:
                            f.write("\tTRAIN_PART_LOSS\tTRAIN_PART_" +
                                    "\tTRAIN_PART_".join(
                                        train_part_eval_result.log_header.
                                        split("\t")))
                        if log_dev:
                            f.write("\tDEV_LOSS\tDEV_" + "\tDEV_".join(
                                dev_eval_result.log_header.split("\t")))
                        if log_test:
                            f.write("\tTEST_LOSS\tTEST_" + "\tTEST_".join(
                                test_eval_result.log_header.split("\t")))

                    f.write(
                        f"\n{self.epoch}\t{datetime.datetime.now():%H:%M:%S}\t{bad_epochs}\t{learning_rate:.4f}\t{train_loss}"
                    )
                    f.write(result_line)

                # if checkpoint is enabled, save model at each epoch
                if checkpoint and not param_selection_mode:
                    self.save_checkpoint(base_path / "checkpoint.pt")

                # if we use dev data, remember best model based on dev evaluation score
                if ((not train_with_dev or anneal_with_restarts
                     or anneal_with_prestarts) and not param_selection_mode
                        and not isinstance(lr_scheduler, OneCycleLR)
                        and current_score == lr_scheduler.best
                        and bad_epochs == 0):
                    print("saving best model")
                    self.model.save(base_path / "best-model.pt")

                    if anneal_with_prestarts:
                        current_state_dict = self.model.state_dict()
                        self.model.load_state_dict(last_epoch_model_state_dict)
                        self.model.save(base_path / "pre-best-model.pt")
                        self.model.load_state_dict(current_state_dict)

                if save_model_at_each_epoch:
                    print("saving model of current epoch")
                    model_name = "model_epoch_" + str(self.epoch) + ".pt"
                    self.model.save(base_path / model_name)

            # if we do not use dev data for model selection, save final model
            if save_final_model and not param_selection_mode:
                self.model.save(base_path / "final-model.pt")

        except KeyboardInterrupt:
            log_line(log)
            log.info("Exiting from training early.")

            if self.use_tensorboard:
                writer.close()

            if not param_selection_mode:
                log.info("Saving model ...")
                self.model.save(base_path / "final-model.pt")
                log.info("Done.")

        # test best model if test data is present
        if self.corpus.test and not train_with_test:
            final_score = self.final_test(base_path, mini_batch_chunk_size,
                                          num_workers)
        else:
            final_score = 0
            log.info("Test data not provided setting final score to 0")

        log.removeHandler(log_handler)

        if self.use_tensorboard:
            writer.close()

        return {
            "test_score": final_score,
            "dev_score_history": dev_score_history,
            "train_loss_history": train_loss_history,
            "dev_loss_history": dev_loss_history,
        }
Esempio n. 23
0
    def predict(
        self,
        sentences: Union[List[Sentence], Sentence],
        mini_batch_size: int = 32,
        return_probabilities_for_all_classes: bool = False,
        verbose: bool = False,
        label_name: Optional[str] = None,
        return_loss=False,
        embedding_storage_mode="none",
    ):
        """
        Predicts the class labels for the given sentences. The labels are directly added to the sentences.
        :param sentences: list of sentences
        :param mini_batch_size: mini batch size to use
        :param return_probabilities_for_all_classes : return probabilities for all classes instead of only best predicted
        :param verbose: set to True to display a progress bar
        :param return_loss: set to True to return loss
        :param label_name: set this to change the name of the label type that is predicted
        :param embedding_storage_mode: default is 'none' which is always best. Only set to 'cpu' or 'gpu' if
        you wish to not only predict, but also keep the generated embeddings in CPU or GPU memory respectively.
        'gpu' to store embeddings in GPU memory.
        """
        if label_name is None:
            label_name = self.label_type if self.label_type is not None else "label"

        with torch.no_grad():
            if not sentences:
                return sentences

            if isinstance(sentences, DataPoint):
                sentences = [sentences]

            # filter empty sentences
            if isinstance(sentences[0], DataPoint):
                sentences = [
                    sentence for sentence in sentences if len(sentence) > 0
                ]
            if len(sentences) == 0:
                return sentences

            # reverse sort all sequences by their length
            rev_order_len_index = sorted(range(len(sentences)),
                                         key=lambda k: len(sentences[k]),
                                         reverse=True)

            reordered_sentences: List[Union[DataPoint, str]] = [
                sentences[index] for index in rev_order_len_index
            ]

            dataloader = DataLoader(
                dataset=SentenceDataset(reordered_sentences),
                batch_size=mini_batch_size)
            # progress bar for verbosity
            if verbose:
                dataloader = tqdm(dataloader)

            overall_loss = 0
            batch_no = 0
            label_count = 0
            for batch in dataloader:

                batch_no += 1

                if verbose:
                    dataloader.set_description(
                        f"Inferencing on batch {batch_no}")

                # stop if all sentences are empty
                if not batch:
                    continue

                scores, gold_labels, data_points, label_candidates = self.forward_pass(
                    batch, return_label_candidates=True)
                # remove previously predicted labels of this type
                for sentence in data_points:
                    sentence.remove_labels(label_name)

                if return_loss:
                    overall_loss += self._calculate_loss(scores,
                                                         gold_labels)[0]
                    label_count += len(label_candidates)

                # if anything could possibly be predicted
                if len(label_candidates) > 0:
                    if self.multi_label:
                        sigmoided = torch.sigmoid(
                            scores)  # size: (n_sentences, n_classes)
                        n_labels = sigmoided.size(1)
                        for s_idx, (data_point, label_candidate) in enumerate(
                                zip(data_points, label_candidates)):
                            for l_idx in range(n_labels):
                                label_value = self.label_dictionary.get_item_for_index(
                                    l_idx)
                                if label_value == 'O': continue
                                label_threshold = self._get_label_threshold(
                                    label_value)
                                label_score = sigmoided[s_idx, l_idx].item()
                                if label_score > label_threshold or return_probabilities_for_all_classes:
                                    label = label_candidate.spawn(
                                        value=label_value, score=label_score)
                                    data_point.add_complex_label(
                                        label_name, label)
                    else:
                        softmax = torch.nn.functional.softmax(scores, dim=-1)

                        if return_probabilities_for_all_classes:
                            n_labels = softmax.size(1)
                            for s_idx, (data_point,
                                        label_candidate) in enumerate(
                                            zip(data_points,
                                                label_candidates)):
                                for l_idx in range(n_labels):
                                    label_value = self.label_dictionary.get_item_for_index(
                                        l_idx)
                                    if label_value == 'O': continue
                                    label_score = softmax[s_idx, l_idx].item()
                                    label = label_candidate.spawn(
                                        value=label_value, score=label_score)
                                    data_point.add_complex_label(
                                        label_name, label)
                        else:
                            conf, idx = torch.max(softmax, dim=-1)
                            for data_point, label_candidate, c, i in zip(
                                    data_points, label_candidates, conf, idx):
                                label_value = self.label_dictionary.get_item_for_index(
                                    i.item())
                                if label_value == 'O': continue
                                label = label_candidate.spawn(
                                    value=label_value, score=c.item())
                                data_point.add_complex_label(label_name, label)

                store_embeddings(batch, storage_mode=embedding_storage_mode)

            if return_loss:
                return overall_loss, label_count
    def evaluate(
        self,
        data_points: Union[List[DataPoint], Dataset],
        gold_label_type: str,
        out_path: Union[str, Path] = None,
        embedding_storage_mode: str = "none",
        mini_batch_size: int = 32,
        num_workers: Optional[int] = 8,
        main_evaluation_metric: Tuple[str, str] = ("micro avg", "f1-score"),
        exclude_labels: List[str] = [],
        gold_label_dictionary: Optional[Dictionary] = None,
        **kwargs,
    ) -> Result:

        if not isinstance(data_points, Dataset):
            data_points = FlairDatapointDataset(data_points)
        data_loader = DataLoader(data_points,
                                 batch_size=mini_batch_size,
                                 num_workers=num_workers)

        lines: List[str] = [
            "token gold_tag gold_arc predicted_tag predicted_arc\n"
        ]

        average_over = 0
        eval_loss_arc = 0.0
        eval_loss_rel = 0.0

        y_true = []
        y_pred = []

        parsing_metric = ParsingMetric()

        for batch in data_loader:
            average_over += 1
            with torch.no_grad():
                score_arc, score_rel = self.forward(batch)
                loss_arc, loss_rel = self._calculate_loss(
                    score_arc, score_rel, batch)
                arc_prediction, relation_prediction = self._obtain_labels_(
                    score_arc, score_rel)

            parsing_metric(arc_prediction, relation_prediction, batch,
                           gold_label_type)

            eval_loss_arc += loss_arc.item()
            eval_loss_rel += loss_rel.item()

            for (sentence, arcs, sent_tags) in zip(batch, arc_prediction,
                                                   relation_prediction):
                for (token, arc, tag) in zip(sentence.tokens, arcs, sent_tags):
                    token.add_tag_label("predicted", Label(tag))
                    token.add_tag_label("predicted_head_id",
                                        Label(str(int(arc))))

                    # append both to file for evaluation
                    eval_line = "{} {} {} {} {}\n".format(
                        token.text,
                        token.get_tag(gold_label_type).value,
                        str(token.head_id),
                        tag,
                        str(int(arc)),
                    )
                    lines.append(eval_line)
                lines.append("\n")

            for sentence in batch:
                gold_tags = [
                    token.get_tag(gold_label_type).value
                    for token in sentence.tokens
                ]
                predicted_tags = [
                    tag.tag for tag in sentence.get_spans("predicted")
                ]

                y_pred += [
                    self.relations_dictionary.get_idx_for_item(tag)
                    for tag in predicted_tags
                ]
                y_true += [
                    self.relations_dictionary.get_idx_for_item(tag)
                    for tag in gold_tags
                ]

            store_embeddings(batch, embedding_storage_mode)

        eval_loss_arc /= average_over
        eval_loss_rel /= average_over

        if out_path is not None:
            with open(out_path, "w", encoding="utf-8") as outfile:
                outfile.write("".join(lines))

        classification_report_dict = sklearn.metrics.classification_report(
            y_true,
            y_pred,
            target_names=self.relations_dictionary.idx2item,
            zero_division=0,
            output_dict=True,
            labels=range(len(self.relations_dictionary)),
        )

        accuracy_score = round(sklearn.metrics.accuracy_score(y_true, y_pred),
                               4)

        precision_score = round(
            classification_report_dict["micro avg"]["precision"], 4)
        recall_score = round(classification_report_dict["micro avg"]["recall"],
                             4)
        micro_f_score = round(
            classification_report_dict["micro avg"]["f1-score"], 4)
        macro_f_score = round(
            classification_report_dict["macro avg"]["f1-score"], 4)

        main_score = classification_report_dict[main_evaluation_metric[0]][
            main_evaluation_metric[1]]

        detailed_result = (
            f"\nUAS : {parsing_metric.get_uas():.4f} - LAS : {parsing_metric.get_las():.4f}"
            f"\neval loss rel : {eval_loss_rel:.4f} - eval loss arc : {eval_loss_arc:.4f}"
            f"\nF-Score: micro : {micro_f_score} - macro : {macro_f_score}"
            f"\n Accuracy: {accuracy_score} - Precision {precision_score} - Recall {recall_score}"
        )
        log_header = "PRECISION\tRECALL\tF1\tACCURACY"
        log_line = f"{precision_score}\t" f"{recall_score}\t" f"{micro_f_score}\t" f"{accuracy_score}"

        result = Result(
            main_score=main_score,
            log_line=log_line,
            log_header=log_header,
            detailed_results=detailed_result,
            classification_report=classification_report_dict,
            loss=(eval_loss_rel + eval_loss_arc),
        )
        return result
Esempio n. 25
0
    def predict(
        self,
        sentences: Union[List[Sentence], Sentence],
        mini_batch_size: int = 32,
        return_probabilities_for_all_classes: bool = False,
        verbose: bool = False,
        label_name: Optional[str] = None,
        return_loss=False,
        embedding_storage_mode="none",
    ):
        """
        Predicts labels for current batch with CRF or Softmax.
        :param sentences: List of sentences in batch
        :param mini_batch_size: batch size for test data
        :param return_probabilities_for_all_classes: Whether to return probabilites for all classes
        :param verbose: whether to use progress bar
        :param label_name: which label to predict
        :param return_loss: whether to return loss value
        :param embedding_storage_mode: determines where to store embeddings - can be "gpu", "cpu" or None.
        """
        if label_name is None:
            label_name = self.tag_type

        with torch.no_grad():
            if not sentences:
                return sentences

            # make sure its a list
            if not isinstance(sentences, list) and not isinstance(
                    sentences, flair.data.Dataset):
                sentences = [sentences]

            # filter empty sentences
            sentences = [
                sentence for sentence in sentences if len(sentence) > 0
            ]

            # reverse sort all sequences by their length
            reordered_sentences = sorted(sentences,
                                         key=lambda s: len(s),
                                         reverse=True)

            if len(reordered_sentences) == 0:
                return sentences

            dataloader = DataLoader(
                dataset=FlairDatapointDataset(reordered_sentences),
                batch_size=mini_batch_size,
            )
            # progress bar for verbosity
            if verbose:
                dataloader = tqdm(dataloader, desc="Batch inference")

            overall_loss = torch.zeros(1, device=flair.device)
            batch_no = 0
            label_count = 0
            for batch in dataloader:

                batch_no += 1

                # stop if all sentences are empty
                if not batch:
                    continue

                # get features from forward propagation
                features, gold_labels = self.forward(batch)

                # remove previously predicted labels of this type
                for sentence in batch:
                    sentence.remove_labels(label_name)

                # if return_loss, get loss value
                if return_loss:
                    loss = self._calculate_loss(features, gold_labels)
                    overall_loss += loss[0]
                    label_count += loss[1]

                # Sort batch in same way as forward propagation
                lengths = torch.LongTensor(
                    [len(sentence) for sentence in batch])
                _, sort_indices = lengths.sort(dim=0, descending=True)
                batch = [batch[i] for i in sort_indices]

                # make predictions
                if self.use_crf:
                    predictions, all_tags = self.viterbi_decoder.decode(
                        features, return_probabilities_for_all_classes)
                else:
                    predictions, all_tags = self._standard_inference(
                        features, batch, return_probabilities_for_all_classes)

                # add predictions to Sentence
                for sentence, sentence_predictions in zip(batch, predictions):

                    # BIOES-labels need to be converted to spans
                    if self.predict_spans:
                        sentence_tags = [
                            label[0] for label in sentence_predictions
                        ]
                        sentence_scores = [
                            label[1] for label in sentence_predictions
                        ]
                        predicted_spans = get_spans_from_bio(
                            sentence_tags, sentence_scores)
                        for predicted_span in predicted_spans:
                            span: Span = sentence[
                                predicted_span[0][0]:predicted_span[0][-1] + 1]
                            span.add_label(label_name,
                                           value=predicted_span[2],
                                           score=predicted_span[1])

                    # token-labels can be added directly
                    else:
                        for token, label in zip(sentence.tokens,
                                                sentence_predictions):
                            token.add_label(typename=label_name,
                                            value=label[0],
                                            score=label[1])

                # all_tags will be empty if all_tag_prob is set to False, so the for loop will be avoided
                for (sentence, sent_all_tags) in zip(batch, all_tags):
                    for (token, token_all_tags) in zip(sentence.tokens,
                                                       sent_all_tags):
                        token.add_tags_proba_dist(label_name, token_all_tags)

                store_embeddings(sentences,
                                 storage_mode=embedding_storage_mode)

            if return_loss:
                return overall_loss, label_count
Esempio n. 26
0
    def train(
            self,
            base_path: Union[Path, str],
            learning_rate: float = 0.1,
            mini_batch_size: int = 32,
            mini_batch_chunk_size: Optional[int] = None,
            max_epochs: int = 100,
            train_with_dev: bool = False,
            train_with_test: bool = False,
            monitor_train: bool = False,
            monitor_test: bool = False,
            main_evaluation_metric: Tuple[str, str] = ("micro avg", 'f1-score'),
            scheduler=AnnealOnPlateau,
            anneal_factor: float = 0.5,
            patience: int = 3,
            min_learning_rate: float = 0.0001,
            initial_extra_patience: int = 0,
            optimizer: torch.optim.Optimizer = SGD,
            cycle_momentum: bool = False,
            warmup_fraction: float = 0.1,
            embeddings_storage_mode: str = "cpu",
            checkpoint: bool = False,
            save_final_model: bool = True,
            anneal_with_restarts: bool = False,
            anneal_with_prestarts: bool = False,
            anneal_against_dev_loss: bool = False,
            batch_growth_annealing: bool = False,
            shuffle: bool = True,
            param_selection_mode: bool = False,
            write_weights: bool = False,
            num_workers: int = 6,
            sampler=None,
            use_amp: bool = False,
            amp_opt_level: str = "O1",
            eval_on_train_fraction: float = 0.0,
            eval_on_train_shuffle: bool = False,
            save_model_each_k_epochs: int = 0,
            tensorboard_comment: str = '',
            use_swa: bool = False,
            use_final_model_for_eval: bool = False,
            gold_label_dictionary_for_eval: Optional[Dictionary] = None,
            create_file_logs: bool = True,
            create_loss_file: bool = True,
            epoch: int = 0,
            use_tensorboard: bool = False,
            tensorboard_log_dir=None,
            metrics_for_tensorboard=[],
            optimizer_state_dict: Optional = None,
            scheduler_state_dict: Optional = None,
            save_optimizer_state: bool = False,
            **kwargs,
    ) -> dict:
        """
        Trains any class that implements the flair.nn.Model interface.
        :param base_path: Main path to which all output during training is logged and models are saved
        :param learning_rate: Initial learning rate (or max, if scheduler is OneCycleLR)
        :param mini_batch_size: Size of mini-batches during training
        :param mini_batch_chunk_size: If mini-batches are larger than this number, they get broken down into chunks of this size for processing purposes
        :param max_epochs: Maximum number of epochs to train. Terminates training if this number is surpassed.
        :param scheduler: The learning rate scheduler to use
        :param checkpoint: If True, a full checkpoint is saved at end of each epoch
        :param cycle_momentum: If scheduler is OneCycleLR, whether the scheduler should cycle also the momentum
        :param anneal_factor: The factor by which the learning rate is annealed
        :param patience: Patience is the number of epochs with no improvement the Trainer waits
         until annealing the learning rate
        :param min_learning_rate: If the learning rate falls below this threshold, training terminates
        :param warmup_fraction: Fraction of warmup steps if the scheduler is LinearSchedulerWithWarmup
        :param train_with_dev:  If True, the data from dev split is added to the training data
        :param train_with_test: If True, the data from test split is added to the training data
        :param monitor_train: If True, training data is evaluated at end of each epoch
        :param monitor_test: If True, test data is evaluated at end of each epoch
        :param embeddings_storage_mode: One of 'none' (all embeddings are deleted and freshly recomputed),
        'cpu' (embeddings are stored on CPU) or 'gpu' (embeddings are stored on GPU)
        :param save_final_model: If True, final model is saved
        :param anneal_with_restarts: If True, the last best model is restored when annealing the learning rate
        :param shuffle: If True, data is shuffled during training
        :param param_selection_mode: If True, testing is performed against dev data. Use this mode when doing
        parameter selection.
        :param num_workers: Number of workers in your data loader.
        :param sampler: You can pass a data sampler here for special sampling of data.
        :param eval_on_train_fraction: the fraction of train data to do the evaluation on,
        if 0. the evaluation is not performed on fraction of training data,
        if 'dev' the size is determined from dev set size
        :param eval_on_train_shuffle: if True the train data fraction is determined on the start of training
        and kept fixed during training, otherwise it's sampled at beginning of each epoch
        :param save_model_each_k_epochs: Each k epochs, a model state will be written out. If set to '5', a model will
        be saved each 5 epochs. Default is 0 which means no model saving.
        :param main_evaluation_metric: Type of metric to use for best model tracking and learning rate scheduling (if dev data is available, otherwise loss will be used), currently only applicable for text_classification_model
        :param tensorboard_comment: Comment to use for tensorboard logging
        :param create_file_logs: If True, the logs will also be stored in a file 'training.log' in the model folder
        :param create_loss_file: If True, the loss will be writen to a file 'loss.tsv' in the model folder
        :param optimizer: The optimizer to use (typically SGD or Adam)
        :param epoch: The starting epoch (normally 0 but could be higher if you continue training model)
        :param use_tensorboard: If True, writes out tensorboard information
        :param tensorboard_log_dir: Directory into which tensorboard log files will be written
        :param metrics_for_tensorboard: List of tuples that specify which metrics (in addition to the main_score) shall be plotted in tensorboard, could be [("macro avg", 'f1-score'), ("macro avg", 'precision')] for example
        :param kwargs: Other arguments for the Optimizer
        :return:
        """

        # create a model card for this model with Flair and PyTorch version
        model_card = {'flair_version': flair.__version__, 'pytorch_version': torch.__version__}

        # also record Transformers version if library is loaded
        try:
            import transformers
            model_card['transformers_version'] = transformers.__version__
        except:
            pass

        # remember all parameters used in train() call
        local_variables = locals()
        training_parameters = {}
        for parameter in signature(self.train).parameters:
            training_parameters[parameter] = local_variables[parameter]
        model_card['training_parameters'] = training_parameters

        # add model card to model
        self.model.model_card = model_card

        if use_tensorboard:
            try:
                from torch.utils.tensorboard import SummaryWriter

                if tensorboard_log_dir is not None and not os.path.exists(tensorboard_log_dir):
                    os.mkdir(tensorboard_log_dir)
                writer = SummaryWriter(log_dir=tensorboard_log_dir, comment=tensorboard_comment)
                log.info(f"tensorboard logging path is {tensorboard_log_dir}")

            except:
                log_line(log)
                log.warning("ATTENTION! PyTorch >= 1.1.0 and pillow are required for TensorBoard support!")
                log_line(log)
                use_tensorboard = False
                pass

        if use_amp:
            if sys.version_info < (3, 0):
                raise RuntimeError("Apex currently only supports Python 3. Aborting.")
            if amp is None:
                raise RuntimeError(
                    "Failed to import apex. Please install apex from https://www.github.com/nvidia/apex "
                    "to enable mixed-precision training."
                )

        if mini_batch_chunk_size is None:
            mini_batch_chunk_size = mini_batch_size
        if learning_rate < min_learning_rate:
            min_learning_rate = learning_rate / 10

        initial_learning_rate = learning_rate

        # cast string to Path
        if type(base_path) is str:
            base_path = Path(base_path)
        base_path.mkdir(exist_ok=True, parents=True)

        if create_file_logs:
            log_handler = add_file_handler(log, base_path / "training.log")
        else:
            log_handler = None

        log_line(log)
        log.info(f'Model: "{self.model}"')
        log_line(log)
        log.info(f'Corpus: "{self.corpus}"')
        log_line(log)
        log.info("Parameters:")
        log.info(f' - learning_rate: "{learning_rate}"')
        log.info(f' - mini_batch_size: "{mini_batch_size}"')
        log.info(f' - patience: "{patience}"')
        log.info(f' - anneal_factor: "{anneal_factor}"')
        log.info(f' - max_epochs: "{max_epochs}"')
        log.info(f' - shuffle: "{shuffle}"')
        log.info(f' - train_with_dev: "{train_with_dev}"')
        log.info(f' - batch_growth_annealing: "{batch_growth_annealing}"')
        log_line(log)
        log.info(f'Model training base path: "{base_path}"')
        log_line(log)
        log.info(f"Device: {flair.device}")
        log_line(log)
        log.info(f"Embeddings storage mode: {embeddings_storage_mode}")
        if isinstance(self.model, SequenceTagger) and self.model.weight_dict and self.model.use_crf:
            log_line(log)
            log.warning(f'WARNING: Specified class weights will not take effect when using CRF')

        # check for previously saved best models in the current training folder and delete them
        self.check_for_and_delete_previous_best_models(base_path)

        # determine what splits (train, dev, test) to evaluate and log
        log_train = True if monitor_train else False
        log_test = True if (not param_selection_mode and self.corpus.test and monitor_test) else False
        log_dev = False if train_with_dev or not self.corpus.dev else True
        log_train_part = True if (eval_on_train_fraction == "dev" or eval_on_train_fraction > 0.0) else False

        if log_train_part:
            train_part_size = len(self.corpus.dev) if eval_on_train_fraction == "dev" \
                else int(len(self.corpus.train) * eval_on_train_fraction)

            assert train_part_size > 0
            if not eval_on_train_shuffle:
                train_part_indices = list(range(train_part_size))
                train_part = torch.utils.data.dataset.Subset(self.corpus.train, train_part_indices)

        # prepare loss logging file and set up header
        loss_txt = init_output_file(base_path, "loss.tsv") if create_loss_file else None

        weight_extractor = WeightExtractor(base_path)

        # if optimizer class is passed, instantiate:
        if inspect.isclass(optimizer):
            optimizer: torch.optim.Optimizer = optimizer(self.model.parameters(), lr=learning_rate, **kwargs)

        if use_swa:
            import torchcontrib
            optimizer = torchcontrib.optim.SWA(optimizer, swa_start=10, swa_freq=5, swa_lr=learning_rate)

        if use_amp:
            self.model, optimizer = amp.initialize(
                self.model, optimizer, opt_level=amp_opt_level
            )

        # load existing optimizer state dictionary if it exists
        if optimizer_state_dict:
            optimizer.load_state_dict(optimizer_state_dict)

        # minimize training loss if training with dev data, else maximize dev score
        anneal_mode = "min" if train_with_dev or anneal_against_dev_loss else "max"
        best_validation_score = 100000000000 if train_with_dev or anneal_against_dev_loss else 0.

        dataset_size = len(self.corpus.train)
        if train_with_dev:
            dataset_size += len(self.corpus.dev)

        # if scheduler is passed as a class, instantiate
        if inspect.isclass(scheduler):
            if scheduler == OneCycleLR:
                scheduler = OneCycleLR(optimizer,
                                       max_lr=learning_rate,
                                       steps_per_epoch=dataset_size // mini_batch_size + 1,
                                       epochs=max_epochs - epoch,
                                       # if we load a checkpoint, we have already trained for epoch
                                       pct_start=0.0,
                                       cycle_momentum=cycle_momentum)
            elif scheduler == LinearSchedulerWithWarmup:
                steps_per_epoch = (dataset_size + mini_batch_size - 1) / mini_batch_size
                num_train_steps = int(steps_per_epoch * max_epochs)
                num_warmup_steps = int(num_train_steps * warmup_fraction)

                scheduler = LinearSchedulerWithWarmup(optimizer,
                                                      num_train_steps=num_train_steps,
                                                      num_warmup_steps=num_warmup_steps)
            else:
                scheduler = scheduler(
                    optimizer,
                    factor=anneal_factor,
                    patience=patience,
                    initial_extra_patience=initial_extra_patience,
                    mode=anneal_mode,
                    verbose=True,
                )

        # load existing scheduler state dictionary if it exists
        if scheduler_state_dict:
            scheduler.load_state_dict(scheduler_state_dict)

        # update optimizer and scheduler in model card
        model_card['training_parameters']['optimizer'] = optimizer
        model_card['training_parameters']['scheduler'] = scheduler

        if isinstance(scheduler, OneCycleLR) and batch_growth_annealing:
            raise ValueError("Batch growth with OneCycle policy is not implemented.")

        train_data = self.corpus.train

        # if training also uses dev/train data, include in training set
        if train_with_dev or train_with_test:

            parts = [self.corpus.train]
            if train_with_dev: parts.append(self.corpus.dev)
            if train_with_test: parts.append(self.corpus.test)

            train_data = ConcatDataset(parts)

        # initialize sampler if provided
        if sampler is not None:
            # init with default values if only class is provided
            if inspect.isclass(sampler):
                sampler = sampler()
            # set dataset to sample from
            sampler.set_dataset(train_data)
            shuffle = False

        dev_score_history = []
        dev_loss_history = []
        train_loss_history = []

        micro_batch_size = mini_batch_chunk_size

        # At any point you can hit Ctrl + C to break out of training early.
        try:
            previous_learning_rate = learning_rate
            momentum = 0
            for group in optimizer.param_groups:
                if "momentum" in group:
                    momentum = group["momentum"]

            for epoch in range(epoch + 1, max_epochs + 1):
                log_line(log)

                # update epoch in model card
                self.model.model_card['training_parameters']['epoch'] = epoch

                if anneal_with_prestarts:
                    last_epoch_model_state_dict = copy.deepcopy(self.model.state_dict())

                if eval_on_train_shuffle:
                    train_part_indices = list(range(self.corpus.train))
                    random.shuffle(train_part_indices)
                    train_part_indices = train_part_indices[:train_part_size]
                    train_part = torch.utils.data.dataset.Subset(self.corpus.train, train_part_indices)

                # get new learning rate
                for group in optimizer.param_groups:
                    learning_rate = group["lr"]

                if learning_rate != previous_learning_rate and batch_growth_annealing:
                    mini_batch_size *= 2

                # reload last best model if annealing with restarts is enabled
                if (
                        (anneal_with_restarts or anneal_with_prestarts)
                        and learning_rate != previous_learning_rate
                        and os.path.exists(base_path / "best-model.pt")
                ):
                    if anneal_with_restarts:
                        log.info("resetting to best model")
                        self.model.load_state_dict(
                            self.model.load(base_path / "best-model.pt").state_dict()
                        )
                    if anneal_with_prestarts:
                        log.info("resetting to pre-best model")
                        self.model.load_state_dict(
                            self.model.load(base_path / "pre-best-model.pt").state_dict()
                        )

                previous_learning_rate = learning_rate
                if use_tensorboard:
                    writer.add_scalar("learning_rate", learning_rate, epoch)

                # stop training if learning rate becomes too small
                if ((not isinstance(scheduler, (OneCycleLR, LinearSchedulerWithWarmup)) and
                     learning_rate < min_learning_rate)):
                    log_line(log)
                    log.info("learning rate too small - quitting training!")
                    log_line(log)
                    break

                batch_loader = DataLoader(
                    train_data,
                    batch_size=mini_batch_size,
                    shuffle=shuffle if epoch > 1 else False,  # never shuffle the first epoch
                    num_workers=num_workers,
                    sampler=sampler,
                )

                self.model.train()

                train_loss: float = 0

                seen_batches = 0
                total_number_of_batches = len(batch_loader)

                modulo = max(1, int(total_number_of_batches / 10))

                # process mini-batches
                batch_time = 0
                average_over = 0
                for batch_no, batch in enumerate(batch_loader):

                    start_time = time.time()

                    # zero the gradients on the model and optimizer
                    self.model.zero_grad()
                    optimizer.zero_grad()

                    # if necessary, make batch_steps
                    batch_steps = [batch]
                    if len(batch) > micro_batch_size:
                        batch_steps = [batch[x: x + micro_batch_size] for x in range(0, len(batch), micro_batch_size)]

                    # forward and backward for batch
                    for batch_step in batch_steps:

                        # forward pass
                        loss = self.model.forward_loss(batch_step)

                        if isinstance(loss, Tuple):
                            average_over += loss[1]
                            loss = loss[0]

                        # Backward
                        if use_amp:
                            with amp.scale_loss(loss, optimizer) as scaled_loss:
                                scaled_loss.backward()
                        else:
                            loss.backward()
                        train_loss += loss.item()

                    # do the optimizer step
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
                    optimizer.step()

                    # do the scheduler step if one-cycle or linear decay
                    if isinstance(scheduler, (OneCycleLR, LinearSchedulerWithWarmup)):
                        scheduler.step()
                        # get new learning rate
                        for group in optimizer.param_groups:
                            learning_rate = group["lr"]
                            if "momentum" in group:
                                momentum = group["momentum"]
                            if "betas" in group:
                                momentum, _ = group["betas"]

                    seen_batches += 1

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(batch, embeddings_storage_mode)

                    batch_time += time.time() - start_time
                    if seen_batches % modulo == 0:
                        momentum_info = f' - momentum: {momentum:.4f}' if cycle_momentum else ''
                        intermittent_loss = train_loss / average_over if average_over > 0 else train_loss / seen_batches
                        log.info(
                            f"epoch {epoch} - iter {seen_batches}/{total_number_of_batches} - loss "
                            f"{intermittent_loss:.8f} - samples/sec: {mini_batch_size * modulo / batch_time:.2f}"
                            f" - lr: {learning_rate:.6f}{momentum_info}"
                        )
                        batch_time = 0
                        iteration = epoch * total_number_of_batches + batch_no
                        if not param_selection_mode and write_weights:
                            weight_extractor.extract_weights(self.model.state_dict(), iteration)

                if average_over != 0:
                    train_loss /= average_over

                self.model.eval()

                log_line(log)
                log.info(f"EPOCH {epoch} done: loss {train_loss:.4f} - lr {learning_rate:.7f}")

                if use_tensorboard:
                    writer.add_scalar("train_loss", train_loss, epoch)

                # evaluate on train / dev / test split depending on training settings
                result_line: str = ""

                if log_train:
                    train_eval_result = self.model.evaluate(
                        self.corpus.train,
                        gold_label_type=self.model.label_type,
                        mini_batch_size=mini_batch_chunk_size,
                        num_workers=num_workers,
                        embedding_storage_mode=embeddings_storage_mode,
                        main_evaluation_metric=main_evaluation_metric,
                        gold_label_dictionary=gold_label_dictionary_for_eval,
                    )
                    result_line += f"\t{train_eval_result.log_line}"

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.train, embeddings_storage_mode)

                if log_train_part:
                    train_part_eval_result = self.model.evaluate(
                        train_part,
                        gold_label_type=self.model.label_type,
                        mini_batch_size=mini_batch_chunk_size,
                        num_workers=num_workers,
                        embedding_storage_mode=embeddings_storage_mode,
                        main_evaluation_metric=main_evaluation_metric,
                        gold_label_dictionary=gold_label_dictionary_for_eval,
                    )
                    result_line += f"\t{train_part_eval_result.loss}\t{train_part_eval_result.log_line}"

                    log.info(
                        f"TRAIN_SPLIT : loss {train_part_eval_result.loss} - {main_evaluation_metric[1]} ({main_evaluation_metric[0]}) {round(train_part_eval_result.main_score, 4)}"
                    )
                if use_tensorboard:
                    for (metric_class_avg_type, metric_type) in metrics_for_tensorboard:
                        writer.add_scalar(
                            f"train_{metric_class_avg_type}_{metric_type}",
                            train_part_eval_result.classification_report[metric_class_avg_type][metric_type], epoch
                        )

                if log_dev:
                    dev_eval_result = self.model.evaluate(
                        self.corpus.dev,
                        gold_label_type=self.model.label_type,
                        mini_batch_size=mini_batch_chunk_size,
                        num_workers=num_workers,
                        out_path=base_path / "dev.tsv",
                        embedding_storage_mode=embeddings_storage_mode,
                        main_evaluation_metric=main_evaluation_metric,
                        gold_label_dictionary=gold_label_dictionary_for_eval,
                    )
                    result_line += f"\t{dev_eval_result.loss}\t{dev_eval_result.log_line}"
                    log.info(
                        f"DEV : loss {dev_eval_result.loss} - {main_evaluation_metric[1]} ({main_evaluation_metric[0]})  {round(dev_eval_result.main_score, 4)}"
                    )
                    # calculate scores using dev data if available
                    # append dev score to score history
                    dev_score_history.append(dev_eval_result.main_score)
                    dev_loss_history.append(dev_eval_result.loss)

                    dev_score = dev_eval_result.main_score

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.dev, embeddings_storage_mode)

                    if use_tensorboard:
                        writer.add_scalar("dev_loss", dev_eval_result.loss, epoch)
                        writer.add_scalar("dev_score", dev_eval_result.main_score, epoch)
                        for (metric_class_avg_type, metric_type) in metrics_for_tensorboard:
                            writer.add_scalar(
                                f"dev_{metric_class_avg_type}_{metric_type}",
                                dev_eval_result.classification_report[metric_class_avg_type][metric_type], epoch
                            )

                if log_test:
                    test_eval_result = self.model.evaluate(
                        self.corpus.test,
                        gold_label_type=self.model.label_type,
                        mini_batch_size=mini_batch_chunk_size,
                        num_workers=num_workers,
                        out_path=base_path / "test.tsv",
                        embedding_storage_mode=embeddings_storage_mode,
                        main_evaluation_metric=main_evaluation_metric,
                        gold_label_dictionary=gold_label_dictionary_for_eval,
                    )
                    result_line += f"\t{test_eval_result.loss}\t{test_eval_result.log_line}"
                    log.info(
                        f"TEST : loss {test_eval_result.loss} - {main_evaluation_metric[1]} ({main_evaluation_metric[0]})  {round(test_eval_result.main_score, 4)}"
                    )

                    # depending on memory mode, embeddings are moved to CPU, GPU or deleted
                    store_embeddings(self.corpus.test, embeddings_storage_mode)

                    if use_tensorboard:
                        writer.add_scalar("test_loss", test_eval_result.loss, epoch)
                        writer.add_scalar("test_score", test_eval_result.main_score, epoch)
                        for (metric_class_avg_type, metric_type) in metrics_for_tensorboard:
                            writer.add_scalar(
                                f"test_{metric_class_avg_type}_{metric_type}",
                                test_eval_result.classification_report[metric_class_avg_type][metric_type], epoch
                            )

                # determine if this is the best model or if we need to anneal
                current_epoch_has_best_model_so_far = False
                # default mode: anneal against dev score
                if not train_with_dev and not anneal_against_dev_loss:
                    if dev_score > best_validation_score:
                        current_epoch_has_best_model_so_far = True
                        best_validation_score = dev_score

                    if isinstance(scheduler, AnnealOnPlateau):
                        scheduler.step(dev_score, dev_eval_result.loss)

                # alternative: anneal against dev loss
                if not train_with_dev and anneal_against_dev_loss:
                    if dev_eval_result.loss < best_validation_score:
                        current_epoch_has_best_model_so_far = True
                        best_validation_score = dev_eval_result.loss

                    if isinstance(scheduler, AnnealOnPlateau):
                        scheduler.step(dev_eval_result.loss)

                # alternative: anneal against train loss
                if train_with_dev:
                    if train_loss < best_validation_score:
                        current_epoch_has_best_model_so_far = True
                        best_validation_score = train_loss

                    if isinstance(scheduler, AnnealOnPlateau):
                        scheduler.step(train_loss)

                train_loss_history.append(train_loss)

                # determine bad epoch number
                try:
                    bad_epochs = scheduler.num_bad_epochs
                except:
                    bad_epochs = 0
                for group in optimizer.param_groups:
                    new_learning_rate = group["lr"]
                if new_learning_rate != previous_learning_rate:
                    bad_epochs = patience + 1
                    if previous_learning_rate == initial_learning_rate: bad_epochs += initial_extra_patience

                # log bad epochs
                log.info(f"BAD EPOCHS (no improvement): {bad_epochs}")

                if create_loss_file:
                    # output log file
                    with open(loss_txt, "a") as f:

                        # make headers on first epoch
                        if epoch == 1:
                            f.write(f"EPOCH\tTIMESTAMP\tBAD_EPOCHS\tLEARNING_RATE\tTRAIN_LOSS")

                            if log_train:
                                f.write("\tTRAIN_" + "\tTRAIN_".join(train_eval_result.log_header.split("\t")))

                            if log_train_part:
                                f.write("\tTRAIN_PART_LOSS\tTRAIN_PART_" + "\tTRAIN_PART_".join(
                                    train_part_eval_result.log_header.split("\t")))

                            if log_dev:
                                f.write("\tDEV_LOSS\tDEV_" + "\tDEV_".join(dev_eval_result.log_header.split("\t")))

                            if log_test:
                                f.write("\tTEST_LOSS\tTEST_" + "\tTEST_".join(test_eval_result.log_header.split("\t")))

                        f.write(
                            f"\n{epoch}\t{datetime.datetime.now():%H:%M:%S}\t{bad_epochs}\t{learning_rate:.4f}\t{train_loss}"
                        )
                        f.write(result_line)

                # if checkpoint is enabled, save model at each epoch
                if checkpoint and not param_selection_mode:
                    self.model.save(base_path / "checkpoint.pt", checkpoint=True)

                # Check whether to save best model
                if (
                        (not train_with_dev or anneal_with_restarts or anneal_with_prestarts)
                        and not param_selection_mode
                        and current_epoch_has_best_model_so_far
                        and not use_final_model_for_eval
                ):
                    log.info("saving best model")
                    self.model.save(base_path / "best-model.pt", checkpoint=save_optimizer_state)

                    if anneal_with_prestarts:
                        current_state_dict = self.model.state_dict()
                        self.model.load_state_dict(last_epoch_model_state_dict)
                        self.model.save(base_path / "pre-best-model.pt")
                        self.model.load_state_dict(current_state_dict)

                if save_model_each_k_epochs > 0 and not epoch % save_model_each_k_epochs:
                    print("saving model of current epoch")
                    model_name = "model_epoch_" + str(epoch) + ".pt"
                    self.model.save(base_path / model_name, checkpoint=save_optimizer_state)

            if use_swa:
                optimizer.swap_swa_sgd()

            # if we do not use dev data for model selection, save final model
            if save_final_model and not param_selection_mode:
                self.model.save(base_path / "final-model.pt", checkpoint=save_optimizer_state)

        except KeyboardInterrupt:
            log_line(log)
            log.info("Exiting from training early.")

            if use_tensorboard:
                writer.close()

            if not param_selection_mode:
                log.info("Saving model ...")
                self.model.save(base_path / "final-model.pt", checkpoint=save_optimizer_state)
                log.info("Done.")

        # test best model if test data is present
        if self.corpus.test and not train_with_test:
            final_score = self.final_test(
                base_path=base_path,
                eval_mini_batch_size=mini_batch_chunk_size,
                num_workers=num_workers,
                main_evaluation_metric=main_evaluation_metric,
                gold_label_dictionary_for_eval=gold_label_dictionary_for_eval,
            )
        else:
            final_score = 0
            log.info("Test data not provided setting final score to 0")

        if create_file_logs:
            log_handler.close()
            log.removeHandler(log_handler)

        if use_tensorboard:
            writer.close()

        return {
            "test_score": final_score,
            "dev_score_history": dev_score_history,
            "train_loss_history": train_loss_history,
            "dev_loss_history": dev_loss_history,
        }
Esempio n. 27
0
    def find_learning_rate(
        self,
        base_path: Union[Path, str],
        file_name: str = "learning_rate.tsv",
        start_learning_rate: float = 1e-7,
        end_learning_rate: float = 10,
        iterations: int = 100,
        mini_batch_size: int = 32,
        stop_early: bool = True,
        smoothing_factor: float = 0.98,
        **kwargs,
    ) -> Path:
        best_loss = None
        moving_avg_loss = 0

        # cast string to Path
        if type(base_path) is str:
            base_path = Path(base_path)
        learning_rate_tsv = init_output_file(base_path, file_name)

        with open(learning_rate_tsv, "a") as f:
            f.write("ITERATION\tTIMESTAMP\tLEARNING_RATE\tTRAIN_LOSS\n")

        optimizer = self.optimizer(self.model.parameters(),
                                   lr=start_learning_rate,
                                   **kwargs)

        train_data = self.corpus.train

        scheduler = ExpAnnealLR(optimizer, end_learning_rate, iterations)

        model_state = self.model.state_dict()
        self.model.train()

        step = 0
        while step < iterations:
            batch_loader = DataLoader(train_data,
                                      batch_size=mini_batch_size,
                                      shuffle=True)
            for batch in batch_loader:
                step += 1

                # forward pass
                loss = self.model.forward_loss(batch)

                # update optimizer and scheduler
                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
                optimizer.step()
                scheduler.step(step)

                print(scheduler.get_lr())
                learning_rate = scheduler.get_lr()[0]

                loss_item = loss.item()
                if step == 1:
                    best_loss = loss_item
                else:
                    if smoothing_factor > 0:
                        moving_avg_loss = (smoothing_factor * moving_avg_loss +
                                           (1 - smoothing_factor) * loss_item)
                        loss_item = moving_avg_loss / (1 - smoothing_factor**
                                                       (step + 1))
                    if loss_item < best_loss:
                        best_loss = loss

                if step > iterations:
                    break

                if stop_early and (loss_item > 4 * best_loss
                                   or torch.isnan(loss)):
                    log_line(log)
                    log.info("loss diverged - stopping early!")
                    step = iterations
                    break

                with open(str(learning_rate_tsv), "a") as f:
                    f.write(
                        f"{step}\t{datetime.datetime.now():%H:%M:%S}\t{learning_rate}\t{loss_item}\n"
                    )

            self.model.load_state_dict(model_state)
            self.model.to(flair.device)

        log_line(log)
        log.info(f"learning rate finder finished - plot {learning_rate_tsv}")
        log_line(log)

        return Path(learning_rate_tsv)
Esempio n. 28
0
    def find_learning_rate(
            self,
            base_path: Union[Path, str],
            optimizer,
            mini_batch_size: int = 32,
            start_learning_rate: float = 1e-7,
            end_learning_rate: float = 10,
            iterations: int = 1000,
            stop_early: bool = True,
            file_name: str = "learning_rate.tsv",
            **kwargs,
    ) -> Path:
        best_loss = None

        # cast string to Path
        if type(base_path) is str:
            base_path = Path(base_path)
        base_path.mkdir(exist_ok=True, parents=True)
        learning_rate_tsv = init_output_file(base_path, file_name)

        with open(learning_rate_tsv, "a") as f:
            f.write("ITERATION\tTIMESTAMP\tLEARNING_RATE\tTRAIN_LOSS\n")

        optimizer = optimizer(self.model.parameters(), lr=start_learning_rate, **kwargs)

        train_data = self.corpus.train

        scheduler = ExpAnnealLR(optimizer, end_learning_rate, iterations)

        model_state = self.model.state_dict()
        self.model.train()

        step = 0

        loss_list = []
        average_loss_list = []

        while step < iterations:

            batch_loader = DataLoader(train_data, batch_size=mini_batch_size, shuffle=True)

            for batch in batch_loader:
                step += 1

                # forward pass
                loss = self.model.forward_loss(batch)
                if isinstance(loss, Tuple):
                    loss = loss[0]

                # update optimizer and scheduler
                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
                optimizer.step()
                scheduler.step()

                learning_rate = scheduler.get_lr()[0]

                # append current loss to list of losses for all iterations
                loss_list.append(loss.item())

                # compute averaged loss
                import statistics
                moving_avg_loss = statistics.mean(loss_list)
                average_loss_list.append(moving_avg_loss)

                if len(average_loss_list) > 10:
                    drop = average_loss_list[-10] - moving_avg_loss
                else:
                    drop = 0.

                if not best_loss or moving_avg_loss < best_loss:
                    best_loss = moving_avg_loss

                if step > iterations:
                    break

                if stop_early and (moving_avg_loss > 4 * best_loss or torch.isnan(loss)):
                    log_line(log)
                    log.info("loss diverged - stopping early!")
                    step = iterations
                    break

                with open(str(learning_rate_tsv), "a") as f:
                    f.write(f"{step}\t{learning_rate}\t{loss.item()}\t{moving_avg_loss}\t{drop}\n")

            self.model.load_state_dict(model_state)
            self.model.to(flair.device)

        log_line(log)
        log.info(f"learning rate finder finished - plot {learning_rate_tsv}")
        log_line(log)

        return Path(learning_rate_tsv)
	def evaluate(
		self,
		data_loader: DataLoader,
		out_path: Path = None,
		embeddings_storage_mode: str = "cpu",
		prediction_mode: bool = False,
	) -> (Result, float):
		data_loader.assign_embeddings()
		with torch.no_grad():
			if self.binary:
				eval_loss = 0

				batch_no: int = 0

				# metric = Metric("Evaluation")
				# sentence_writer = open('temps/'+str(uid)+'_eval'+'.conllu','w')
				lines: List[str] = []
				utp = 0
				ufp = 0
				ufn = 0
				ltp = 0
				lfp = 0
				lfn = 0
				for batch in data_loader:
					batch_no += 1
					
					arc_scores, rel_scores = self.forward(batch)
					mask=self.mask
					root_mask = mask.clone()
					root_mask[:,0] = 0
					binary_mask = root_mask.unsqueeze(-1) * mask.unsqueeze(-2)
					
					arc_predictions = (arc_scores.sigmoid() > 0.5) * binary_mask
					rel_predictions = (rel_scores.softmax(-1)*binary_mask.unsqueeze(-1)).argmax(-1)
					if not prediction_mode:
						arc_mat=torch.stack([getattr(sentence,self.tag_type+'_arc_tags').to(flair.device) for sentence in batch],0).float()
						rel_mat=torch.stack([getattr(sentence,self.tag_type+'_rel_tags').to(flair.device) for sentence in batch],0).long()
						loss = self._calculate_loss(arc_scores, rel_scores, batch, mask)
						if self.is_srl:
							# let the head selection fixed to the gold predicate only
							binary_mask[:,:,0] = arc_mat[:,:,0]
							arc_predictions = (arc_scores.sigmoid() > 0.5) * binary_mask
							

						# UF1
						true_positives = arc_predictions * arc_mat
						# (n x m x m) -> ()
						n_predictions = arc_predictions.sum()
						n_unlabeled_predictions = n_predictions
						n_targets = arc_mat.sum()
						n_unlabeled_targets = n_targets
						n_true_positives = true_positives.sum()
						# () - () -> ()
						n_false_positives = n_predictions - n_true_positives
						n_false_negatives = n_targets - n_true_positives
						# (n x m x m) -> (n)
						n_targets_per_sequence = arc_mat.sum([1,2])
						n_true_positives_per_sequence = true_positives.sum([1,2])
						# (n) x 2 -> ()
						n_correct_sequences = (n_true_positives_per_sequence==n_targets_per_sequence).sum()
						utp += n_true_positives
						ufp += n_false_positives
						ufn += n_false_negatives

						# LF1
						# (n x m x m) (*) (n x m x m) -> (n x m x m)
						true_positives = (rel_predictions == rel_mat) * arc_predictions
						correct_label_tokens = (rel_predictions == rel_mat) * arc_mat
						# (n x m x m) -> ()
						# n_unlabeled_predictions = tf.reduce_sum(unlabeled_predictions)
						# n_unlabeled_targets = tf.reduce_sum(unlabeled_targets)
						n_true_positives = true_positives.sum()
						n_correct_label_tokens = correct_label_tokens.sum()
						# () - () -> ()
						n_false_positives = n_unlabeled_predictions - n_true_positives
						n_false_negatives = n_unlabeled_targets - n_true_positives
						# (n x m x m) -> (n)
						n_targets_per_sequence = arc_mat.sum([1,2])
						n_true_positives_per_sequence = true_positives.sum([1,2])
						n_correct_label_tokens_per_sequence = correct_label_tokens.sum([1,2])
						# (n) x 2 -> ()
						n_correct_sequences = (n_true_positives_per_sequence == n_targets_per_sequence).sum()
						n_correct_label_sequences = ((n_correct_label_tokens_per_sequence == n_targets_per_sequence)).sum()
						ltp += n_true_positives
						lfp += n_false_positives
						lfn += n_false_negatives

						eval_loss += loss
						eval_loss /= batch_no
						UF1=self.compute_F1(utp,ufp,ufn)
						LF1=self.compute_F1(ltp,lfp,lfn)

					if out_path is not None:
						masked_arc_scores = arc_scores.masked_fill(~binary_mask.bool(), float(-1e9))
						# if self.target
						# lengths = [len(sentence.tokens) for sentence in batch]
						
						# temp_preds = eisner(arc_scores, mask)
						if not self.is_mst:
							temp_preds = eisner(arc_scores, root_mask.bool())
						for (sent_idx, sentence) in enumerate(batch):
							if self.is_mst:
								preds=MST_inference(torch.softmax(masked_arc_scores[sent_idx],-1).cpu().numpy(), len(sentence), binary_mask[sent_idx].cpu().numpy())
							else:
								preds=temp_preds[sent_idx]

							for token_idx, token in enumerate(sentence):
								if token_idx == 0:
									continue

								# append both to file for evaluation
								arc_heads = torch.where(arc_predictions[sent_idx,token_idx]>0)[0]
								if preds[token_idx] not in arc_heads:
									val=torch.zeros(1).type_as(arc_heads)
									val[0]=preds[token_idx].item()
									arc_heads=torch.cat([arc_heads,val],0)
								if len(arc_heads) == 0:
									arc_heads = masked_arc_scores[sent_idx,token_idx].argmax().unsqueeze(0)
								rel_index = rel_predictions[sent_idx,token_idx,arc_heads]
								rel_labels = [self.tag_dictionary.get_item_for_index(x) for x in rel_index]
								arc_list=[]
								for i, label in enumerate(rel_labels):
									if '+' in label:
										labels = label.split('+')
										for temp_label in labels:
											arc_list.append(str(arc_heads[i].item())+':'+temp_label)
									else:
										arc_list.append(str(arc_heads[i].item())+':'+label)
								eval_line = "{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\n".format(
									token_idx,
									token.text,
									'X',
									'X',
									'X',
									'X=X',
									str(token_idx-1),
									'root' if token_idx-1==0 else 'det',
									'|'.join(arc_list),
									'X',
								)
								lines.append(eval_line)
							lines.append("\n")

				if out_path is not None:
					with open(out_path, "w", encoding="utf-8") as outfile:
						outfile.write("".join(lines))
				if prediction_mode:
					return None, None

				result = Result(
					main_score=LF1,
					log_line=f"\nUF1: {UF1} - LF1 {LF1}",
					log_header="PRECISION\tRECALL\tF1",
					detailed_results=f"\nUF1: {UF1} - LF1 {LF1}",
				)
			else:
				if prediction_mode:
					eval_loss, metric=self.dependency_evaluate(data_loader,out_path=out_path,prediction_mode=prediction_mode)
					return eval_loss, metric
				else:   
					eval_loss, metric=self.dependency_evaluate(data_loader,out_path=out_path)
				
				UAS=metric.uas
				LAS=metric.las
				result = Result(main_score=LAS,log_line=f"\nUAS: {UAS} - LAS {LAS}",log_header="PRECISION\tRECALL\tF1",detailed_results=f"\nUAS: {UAS} - LAS {LAS}",)
			return result, eval_loss
    def evaluate(self,
                 sentences: Union[List[DataPoint], Dataset],
                 out_path: Union[str, Path] = None,
                 embedding_storage_mode: str = "none",
                 mini_batch_size: int = 32,
                 num_workers: int = 8,
                 **kwargs) -> (Result, float):

        # read Dataset into data loader (if list of sentences passed, make Dataset first)
        if not isinstance(sentences, Dataset):
            sentences = SentenceDataset(sentences)
        data_loader = DataLoader(sentences,
                                 batch_size=mini_batch_size,
                                 num_workers=num_workers)

        with torch.no_grad():
            eval_loss = 0

            metric = MetricRegression("Evaluation")

            lines: List[str] = []
            total_count = 0
            for batch_nr, batch in enumerate(data_loader):

                if isinstance(batch, Sentence):
                    batch = [batch]

                scores, loss = self.forward_labels_and_loss(batch)

                true_values = []
                for sentence in batch:
                    total_count += 1
                    for label in sentence.labels:
                        true_values.append(float(label.value))

                results = []
                for score in scores:
                    if type(score[0]) is Label:
                        results.append(float(score[0].score))
                    else:
                        results.append(float(score[0]))

                eval_loss += loss

                metric.true.extend(true_values)
                metric.pred.extend(results)

                for sentence, prediction, true_value in zip(
                        batch, results, true_values):
                    eval_line = "{}\t{}\t{}\n".format(
                        sentence.to_original_text(), true_value, prediction)
                    lines.append(eval_line)

                store_embeddings(batch, embedding_storage_mode)

            eval_loss /= total_count

            ##TODO: not saving lines yet
            if out_path is not None:
                with open(out_path, "w", encoding="utf-8") as outfile:
                    outfile.write("".join(lines))

            log_line = f"{metric.mean_squared_error()}\t{metric.spearmanr()}\t{metric.pearsonr()}"
            log_header = "MSE\tSPEARMAN\tPEARSON"

            detailed_result = (
                f"AVG: mse: {metric.mean_squared_error():.4f} - "
                f"mae: {metric.mean_absolute_error():.4f} - "
                f"pearson: {metric.pearsonr():.4f} - "
                f"spearman: {metric.spearmanr():.4f}")

            result: Result = Result(
                main_score=metric.pearsonr(),
                loss=eval_loss,
                log_header=log_header,
                log_line=log_line,
                detailed_results=detailed_result,
            )

            return result