Closed RobertTLange closed 3 years ago
Ideally the model logging/storage is even more flexible and also allows for simpler sklearn
style setups. The log_config
dictionary could include a model_type
variable that is in [sklearn, torch, tensorflow, jax]
and then stores/logs the different statistics/checkpoints/model states accordingly.
TIL: sklearn
models/trained classifiers can be simply stored via pickle: https://stackoverflow.com/questions/10592605/save-classifier-to-disk-in-scikit-learn
Addressed in PR #29.
Point 1: Currently we only have two options to store checkpoints:
.pt
file.a) Allow for the option to store top k checkpoints based on a logged variable that measures the performance. b) Is there a better way to index checkpoints in order to figure out when the snapshot was taken? Store that in corresponding log .hdf5 file?
Point 2: Network stats logging right now only works if network is torch based. Let's add an option to differentiate between torch and JAX. E.g. based on whether
network
is atorch.nn
module or aFrozenDict
for JAX.