mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-01 12:32:04 +08:00
William Falcon
7099f8dbfb
* split trainer mixins * Update multi_node_cluster_template.py * Update single_cpu_template.py * Update single_gpu_node_16bit_template.py * Update single_gpu_node_ddp_template.py * Update single_gpu_node_dp_template.py * Update trainer_cpu_template.py * Update trainer_io.py * split trainer mixins * Update multi_node_cluster_template.py * deconflicted * deconflicted * deconflicted
35 lines
890 B
Python
35 lines
890 B
Python
class ModelIO(object):
|
|
|
|
def on_load_checkpoint(self, checkpoint):
|
|
"""
|
|
Do something with the checkpoint
|
|
Gives model a chance to load something before state_dict is restored
|
|
:param checkpoint:
|
|
:return:
|
|
"""
|
|
pass
|
|
|
|
def on_save_checkpoint(self, checkpoint):
|
|
"""
|
|
Give the model a chance to add something to the checkpoint.
|
|
state_dict is already there
|
|
"""
|
|
pass
|
|
|
|
# -------------------------
|
|
# OPTIONAL HOOKS
|
|
# -------------------------
|
|
def on_hpc_save(self, checkpoint):
|
|
"""
|
|
Hook to do whatever you need right before Slurm manager saves the model
|
|
:return:
|
|
"""
|
|
pass
|
|
|
|
def on_hpc_load(self, checkpoint):
|
|
"""
|
|
Hook to do whatever you need right before Slurm manager loads the model
|
|
:return:
|
|
"""
|
|
pass
|