optimizer save fix
@@ -2,15 +2,27 @@ import os
import dill
import numpy as np
import pathlib
+import torch
class SolverConfig():
conditions = None
ode_system = None
+ nets = None
+ optimizer = None
+ optimizer_params = {}
+ train_generator = None
+ valid_generator = None
class PretrainedSolver():
#Saving selected attributes of model in dict
def save(self,filename):
+ # Check if optimizer is existing in pytorch
+ optimizer_class=None
+ for cls in torch.optim.Optimizer.__subclasses__():
+ if self.optimizer.__class__.__name__ == cls.__name__:
+ optimizer_class = self.optimizer.__class__
+
save_dict = {
"metrics": self.metrics_fn,
"criterion": self.criterion,
@@ -18,6 +30,8 @@ class PretrainedSolver():
"global_epoch": self.global_epoch, #loss_history
"nets": self.nets,
"optimizer": self.optimizer,
+ "optimizer_state": self.optimizer.state_dict()
+ "optimizer_class": optimizer_class,
"diff_eqs": self.diff_eqs,
"generator": self.generator,
"type": self.__class__
@@ -27,36 +41,64 @@ class PretrainedSolver():
dill.dump(save_dict,file)
#Loading saved attributes into new solver object
+
+ # Have to add check and warning/error
@classmethod
def load(cls, path, config=SolverConfig()):
with open(path,'rb') as file:
load_dict = dill.load(file)
- t = load_dict['generator']['train'].get_examples()[0].detach().numpy()
+ # Loading user defined generator and extracting time domain information
+ if config.train_generator == None:
+ train_generator = load_dict['generator']['train']
+ valid_generator = load_dict['generator']['train']
+ else:
+ train_generator = config.train_generator
+ valid_generator = config.valid_generator
+
+ t = train_generator.get_examples()[0].detach().numpy()
t_min = np.round(min(t))
t_max = np.round(max(t))
- # For 1D
- # params = {"ode_system":load_dict['diff_eqs'],"key":2}
- # solver = cls(**params)
+ # Loading user defined ode_system or system from load file
if config.ode_system == None:
ode = load_dict['diff_eqs']
else:
ode = config.ode_system
+ # Loading user defined conditions or conditions from load file
if config.conditions == None:
cond = load_dict['conditions']
else:
cond = config.conditions
+ # Loading user defined nets or nets from load file
+ if config.nets == None:
+ nets = load_dict['nets']
+ else:
+ nets = config.nets
+
+ # Loading user defined optimizer or optimizer from load file
+ if config.optimizer == None:
+ if load_dict['optimizer_class'] == None:
+ optimizer = load_dict['optimizer']
+
+ else if config.optimizer_params:
+ optimizer = optimizer.load_state_dict(load_dict['opimizer_params'])
+ # net is new, save optimizer class
+ else if type(config.optimizer) == type:
+ #get net and params
+ optimizer = config.optimizer(,**config.optimizer_params)
+ else :
+ optimizer = config.optimizer
solver = cls(ode_system = ode,
conditions = cond,
criterion = load_dict['criterion'],
metrics = load_dict['metrics'],
- nets = load_dict['nets'],
- optimizer = load_dict['optimizer'],
- train_generator = load_dict['generator']['train'],
- valid_generator = load_dict['generator']['valid'],
+ nets = nets,
+ optimizer = optimizer,
+ train_generator = train_generator,
+ valid_generator = valid_generator,
t_min = t_min,
t_max = t_max)
