How to track the best metric score during training (Pytroch Lightning)

58 views Asked by At

I'm learning Pytorch Lightning.

I want to track loss and performance metric scores of each epoch for saving best model or plot chart or get the highest performance but all I have is the the loss and performance score of a batch in the training_step method like this:


class Model(pl.LightningModule):
    ...

    def training_step(self, batch, batch_idx):
        ...

        probs = self.forward(input_ids, labels, output_ids)
        loss = self.loss(probs, labels)
        
        # calculate blue_score
        predictions = self.decode_prediction(probs, input_tokens)
        output_sequences = [[x] for x in output_sequences]
        blue_score = self.blue_score(predictions, output_sequences)
        
        self.log_dict(
            {'train_loss': loss, 'blue_score': blue_score},
            on_step=True,
            on_epoch=True,
            prog_bar=True
        )

        return loss

Edit:

After doing some researches. I think I figured it out. As the default logger already compute accumulate loss and metric score on each epoch (I set on_epoch = True). Now I will write a custom Logger inherit the Logger class of Torch Lightning:

class HistoryLogger(Logger):
    def __init__(self):
        super().__init__()
        self.history = collections.defaultdict(list)

    @property
    def name(self):
        return "HistoryLogger"

    @property
    def version(self):
        return "1.0"

    @rank_zero_only
    def log_metrics(self, metrics, step):            
        for metric_name, metric_value in metrics.items():
            self.history[metric_name].append(metric_value)
        return

logger = HistoryLogger()

trainer = Trainer(logger=logger)
...

# access the history

logger.history
0

There are 0 answers