Class FinetuneCheckPoint

ppq/quantization/algorithm/training.py:93–125

Finetune Check Point stores training loss for variables. It bounds to a specific variable, collects and stores its fp32 value as a reference. ATTENTION: collecting fp32 value might cause GPU memory overflow, so we use a seed to sample only a part of fp32 value instead(randomly pick

class FinetuneCheckPoint:
    """Finetune Check Point stores training loss for variables. It bounds to a
    specific variable, collects and stores its fp32 value as a reference.

    ATTENTION: collecting fp32 value might cause GPU memory overflow, so we use a seed to
        sample only a part of fp32 value instead(randomly pick about 2000 values from given tensor).

    Finetune Check Point maintains a seed for data collecting, a best loss, and a reference values.
    """
    def __init__(self, variable: str, random_fetch: bool = True, seed: int=None, fetchs: int=NUM_OF_CHECKPOINT_FETCHS) -> None:
        if seed is None: seed = randint(0, 0xffffffff)
        self.monitor_var = variable
        self.best_loss   = float(1e9)
        self.seed        = seed
        self.references  = []
        self.outputs     = []
        self.fetchs      = fetchs
        self.random_fetch = random_fetch

    def push(self, tensor: torch.Tensor, is_reference: bool) -> None:
        if self.random_fetch:
            tensor = batch_random_fetch(tensor, seed=self.seed, fetches_per_batch=self.fetchs)
        if is_reference: self.references.append(tensor)
        else: self.outputs.append(tensor)

    def pop(self) -> Tuple[torch.Tensor]:
        assert len(self.outputs) == len(self.references), ('Inconsistent samples detected.'
            f'Reference output gets {len(self.references)} samples, however output has {len(self.outputs)}.')

        return self.outputs, self.references

    def clear(self):
        self.outputs.clear()


class RandomMemDataset: