mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-23 13:40:39 +08:00
3c2fd560aa
* first pass for LightningModule typehints * fix return types * add missing types * add type annotations to grads.py * add type annotations to hooks.py * add type annotation to memory.py * proper docstring quotation marks * add type annotations to saving.py * fix cyclic import problem * fix cyclic import problem * add missing whitespace * finish type hints for load_from_ methods * docs: prepare_data does not return anything * fix auto types in docs * revert typehint for trainer in hook * remove unnecessary return docs * some fixes for memory docs * revert typing for args kwargs * added all missing None return types * remove unused import * add more details to dict/list return types * fix line too long * optimize imports * linted * Revert "linted" This reverts commit 85559611e84e312bce64f4e73b638d4999a8439e. * remove whitespace * update * update * update * update * update * changelog Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> Co-authored-by: William Falcon <waf2107@columbia.edu>
65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
import csv
|
|
import logging as log
|
|
import os
|
|
from argparse import Namespace
|
|
from typing import Union, Dict, Any
|
|
|
|
|
|
class ModelIO(object):
|
|
|
|
def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
|
|
"""
|
|
Do something with the checkpoint
|
|
Gives model a chance to load something before state_dict is restored
|
|
:param checkpoint:
|
|
:return:
|
|
"""
|
|
|
|
def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
|
|
"""
|
|
Give the model a chance to add something to the checkpoint.
|
|
state_dict is already there
|
|
"""
|
|
|
|
# -------------------------
|
|
# OPTIONAL HOOKS
|
|
# -------------------------
|
|
def on_hpc_save(self, checkpoint: Dict[str, Any]) -> None:
|
|
"""
|
|
Hook to do whatever you need right before Slurm manager saves the model
|
|
"""
|
|
|
|
def on_hpc_load(self, checkpoint: Dict[str, Any]) -> None:
|
|
"""
|
|
Hook to do whatever you need right before Slurm manager loads the model
|
|
"""
|
|
|
|
|
|
def load_hparams_from_tags_csv(tags_csv: str) -> Namespace:
|
|
if not os.path.isfile(tags_csv):
|
|
log.warning(f'Missing Tags: {tags_csv}.')
|
|
return Namespace()
|
|
|
|
with open(tags_csv) as f:
|
|
csv_reader = csv.reader(f, delimiter=',')
|
|
tags = {row[0]: convert(row[1]) for row in list(csv_reader)[1:]}
|
|
ns = Namespace(**tags)
|
|
return ns
|
|
|
|
|
|
def convert(val: str) -> Union[int, float, bool, str]:
|
|
constructors = [int, float, str]
|
|
|
|
if isinstance(val, str):
|
|
if val.lower() == 'true':
|
|
return True
|
|
if val.lower() == 'false':
|
|
return False
|
|
|
|
for c in constructors:
|
|
try:
|
|
return c(val)
|
|
except ValueError:
|
|
pass
|
|
return val
|