flytekit.extras.pytorch.PyTorchCheckpoint

class flytekit.extras.pytorch.PyTorchCheckpoint(module=None, hyperparameters=None, optimizer=None)

This class is helpful to save a checkpoint.

Methods

Parameters

• module (Optional[torch.nn.modules.module.Module]) –

• hyperparameters (Optional[Union[Dict[str, Any], NamedTuple, flytekit.extras.pytorch.checkpoint.IsDataclass]]) –

• optimizer (Optional[torch.optim.optimizer.Optimizer]) –

Return type

None

classmethod from_dict(kvs, *, infer_missing=False)

Parameters

kvs (Optional[Union[dict, list, str, int, float, bool]]) –

Return type

dataclasses_json.api.A

classmethod from_json(s, *, parse_float=None, parse_int=None, parse_constant=None, infer_missing=False, **kw)

Parameters

s (Union[str, bytes, bytearray]) –

Return type

dataclasses_json.api.A

classmethod schema(*, infer_missing=False, only=None, exclude=(), many=False, context=None, load_only=(), dump_only=(), partial=False, unknown=None)

Parameters

• infer_missing (bool) –

• many (bool) –

• partial (bool) –

Return type

dataclasses_json.mm.SchemaF[dataclasses_json.mm.A]

to_dict(encode_json=False)

Return type

Dict[str, Optional[Union[dict, list, str, int, float, bool]]]

to_json(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, indent=None, separators=None, default=None, sort_keys=False, **kw)

Parameters

• skipkeys (bool) –

• ensure_ascii (bool) –

• check_circular (bool) –

• allow_nan (bool) –

• indent (Optional[Union[int, str]]) –

• separators (Optional[Tuple[str, str]]) –

• default (Optional[Callable]) –

• sort_keys (bool) –

Return type

str

Attributes

hyperparameters: Optional[Union[Dict[str, Any], NamedTuple, flytekit.extras.pytorch.checkpoint.IsDataclass]] = None

module: Optional[torch.nn.modules.module.Module] = None

optimizer: Optional[torch.optim.optimizer.Optimizer] = None