Class AppState

trainer/base_trainer.py:315–353

This is a useful wrapper for checkpointing the Application State. Since this object is compliant with the Stateful protocol, DCP will automatically call state_dict/load_stat_dict as needed in the dcp.save/load APIs. Note: We take advantage of this wrapper to handle calling distributed s

                f.write(txt)
            
class AppState(Stateful):
    """This is a useful wrapper for checkpointing the Application State. Since
    this object is compliant with the Stateful protocol, DCP will automatically
    call state_dict/load_stat_dict as needed in the dcp.save/load APIs.

    Note: We take advantage of this wrapper to handle calling distributed state dict methods on the model
    and optimizer.
    """

    def __init__(self, model, optimizer=None):
        self.model = model
        self.optimizer = optimizer

    def state_dict(self):
        # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
        model_state_dict, optimizer_state_dict = get_state_dict(
            self.model, self.optimizer)
        state_dict = {
            'model': model_state_dict,
            'optim': optimizer_state_dict,
        }
        if self.ema is not None:
            ema_state_dict, _ = get_state_dict(self.ema, self.optimizer)
            state_dict['ema'] = ema_state_dict
        return state_dict

    def load_state_dict(self, state_dict):
        # sets our state dicts on the model and optimizer, now that we've loaded
        set_state_dict(
            self.model,
            self.optimizer,
            model_state_dict=state_dict['model'],
            optim_state_dict=state_dict['optim'],
        )
        if self.ema is not None:
            set_state_dict(
                self.ema,
                model_state_dict=state_dict['ema'],
            )


class Timer: