def check_init(config, config_kwargs, prnt, activation=None):
if not config_kwargs:
config_kwargs = {}
assert isinstance(config_kwargs, dict)
settings = config(**config_kwargs)
settings.DEBUG = False
#settings.ACTIVATION = activation
assert isinstance(settings, Config)
model = ML_utils.load_model_from_settings(settings)
print(settings.__class__.__name__)
if config_kwargs != {}:
for k, v in config_kwargs.items():
print("{}: {}".format(
k,
v,
), end=", ")
print(end="\n")
if prnt:
print(model.layers_encode)
num_params = sum(p.numel() for p in model.parameters())
print("num params", num_params)
print("CHANNELS", settings.get_channels())
def DA_AE(self, force_init=False, save_vtu=False):
if self.data.get("model") == None or force_init:
self.model = ML_utils.load_model_from_settings(self.settings, self.data.get("device"))
self.data["model"] = self.model
else:
self.model = self.data.get("model")
self.data["model"].eval()
if self.settings.REDUCED_SPACE:
if self.data.get("V_trunc") is None or force_init: #only init if not already init
V_red = VDAInit.create_V_red(self.data.get("train_X"),
self.data.get("encoder"),
self.settings)
self.data["V_trunc"] = V_red.T #tau x M
self.data["w_0"] = np.zeros((V_red.shape[0]))
if self.data["G"] is 1:
self.data["G_V"] =self.data["V_trunc"]
else:
self.data["G_V"] = (self.data["G"] @ self.data["V_trunc"] ).astype(float)
self.data["V_grad"] = None
else:
# Now access explicit gradient function
self.data["V_grad"] = self.__maybe_get_jacobian()
DA_results = self.perform_VarDA(self.data, self.settings, save_vtu=save_vtu)
return DA_results
def __init__(self,
AE_settings,
expdir,
batch_sz=BATCH,
model=None,
start_epoch=None):
"""Initilaizes the AE training class.
::AE_settings - a settings.config.Config class with the DA settings
::expdir - a directory of form `experiments/<possible_path>` to keep logs
::calc_DA_MAE - boolean. If True, training will evaluate DA Mean Absolute Error
during the training cycle. Note: this is *MUCH* slower
"""
self.settings = AE_settings
err_msg = """AE_settings must be an AE configuration class"""
assert self.settings.COMPRESSION_METHOD == "AE", err_msg
if model is not None: #for retraining
assert start_epoch is not None, "If you are RE-training model you must pass start_epoch"
assert start_epoch >= 0
self.start_epoch = start_epoch
self.model = model
print("Loaded model, ", end="")
else:
self.start_epoch = 0
self.model = ML_utils.load_model_from_settings(AE_settings)
print("Initialized model, ", end="")
print("Number of parameters:",
sum(p.numel() for p in self.model.parameters()))
self.batch_sz = batch_sz
self.settings.batch_sz = batch_sz
self.expdir = init_expdir(expdir)
self.settings_fp = self.expdir + "settings.txt"
if self.settings.SAVE == True:
with open(self.settings_fp, "wb") as f:
pickle.dump(self.settings, f)
ML_utils.set_seeds() #set seeds before init model
self.device = ML_utils.get_device()
self.columns = [
"epoch", "reconstruction_err", "DA_MAE", "DA_ratio_improve_MAE",
"time_DA(s)", "time_epoch(s)"
]
def run(self):
"""Generates matrices for VarDA. All returned matrices are in the
(M X n) or (M x nx x ny x nz) format """
data = {}
loader = self.settings.get_loader()
splitter = SplitData()
settings = self.settings
X = loader.get_X(settings)
train_X, test_X, u_c_std, X, mean, std = splitter.train_test_DA_split_maybe_normalize(
X, settings)
if self.u_c is None:
self.u_c = u_c_std
#self.u_c = train_X[62] #good
#self.u_c = train_X[-1] #bad
# We will take initial condition u_0, as mean of historical data
if settings.NORMALIZE:
u_0 = np.zeros_like(mean) #since the data is mean centred
else:
u_0 = mean
encoder = None
decoder = None
device = ML_utils.get_device()
model = self.AEmodel
if model:
model.to(device)
if self.settings.COMPRESSION_METHOD == "AE":
#get encoder
if model is None:
model = ML_utils.load_model_from_settings(settings)
def __create_encoderOrDecoder(fn):
"""This returns a function that deals with encoder/decoder
input dimensions (e.g. adds channel dim for 3D case)"""
def ret_fn(vec):
vec = torch.Tensor(vec).to(device)
#for 3D case, unsqueeze for channel
if self.settings.THREE_DIM:
dims = len(vec.shape)
if dims == 3:
vec = vec.unsqueeze(0)
elif dims == 4:
#batched input
vec = vec.unsqueeze(1)
with torch.no_grad():
res = fn(vec).detach().cpu()
#for 3D case, squeeze for channel
dims = len(res.shape)
if self.settings.THREE_DIM and dims > 2:
if dims == 4:
res = res.squeeze(0)
elif dims == 5: #batched input
res = res.squeeze(1)
return res.numpy()
return ret_fn
encoder = __create_encoderOrDecoder(model.encode)
decoder = __create_encoderOrDecoder(model.decode)
H_0, obs_idx = None, None
if self.settings.REDUCED_SPACE == True:
if self.settings.COMPRESSION_METHOD == "SVD":
raise NotImplementedError(
"SVD in reduced space not implemented")
self.settings.OBS_MODE = "all"
observations, H_0, w_0, d = self.__get_obs_and_d_reduced_space(
self.settings, self.u_c, u_0, encoder)
else:
observations, w_0, d, obs_idx = self.__get_obs_and_d_not_reduced(
self.settings, self.u_c, u_0, encoder)
#TODO - **maybe** get rid of this monstrosity...:
#i.e. you could return a class that has these attributes:
data = {
"d": d,
"G": H_0,
"observations": observations,
"model": model,
"obs_idx": obs_idx,
"encoder": encoder,
"decoder": decoder,
"u_c": self.u_c,
"u_0": u_0,
"X": X,
"train_X": train_X,
"test_X": test_X,
"std": std,
"mean": mean,
"device": device
}
if w_0 is not None:
data["w_0"] = w_0
return data
def __maybe_cross_val_lr(self, test_every, num_epochs_cv=8):
if not num_epochs_cv:
self.num_epochs_cv = 0
return self.learning_rate
elif self.num_epoch < num_epochs_cv:
self.num_epochs_cv = self.num_epoch
else:
self.num_epochs_cv = num_epochs_cv
mult = 1
if self.settings.BATCH_NORM: #i.e. generally larger learning_rate with BN
mult = 5
mult *= BATCH_MULT #linear multiply by size of batch: https://arxiv.org/abs/1706.02677
lrs_base = [0.0001, 0.0003, 0.001]
lrs = [mult * x for x in lrs_base]
res = []
optimizers = []
for idx, lr in enumerate(lrs):
ML_utils.set_seeds() #set seeds before init model
self.model = ML_utils.load_model_from_settings(self.settings)
self.optimizer = optim.Adam(self.model.parameters(), lr)
test_losses = []
train_losses = []
print("learning rate:", lr)
for epoch in range(self.start_epoch,
self.num_epochs_cv + self.start_epoch):
self.epoch = epoch
train, test = self.train_one_epoch(epoch, self.print_every,
test_every,
self.num_epochs_cv)
if test:
test_losses.append(test)
train_losses.append(train)
df = pd.DataFrame(train_losses, columns=self.columns)
train_final = df.tail(1).reconstruction_err
res.append(train_final.values[0])
optimizers.append(self.optimizer)
#save model if best so far
if res[-1] == min(res):
best_test = test_losses
best_train = train_losses
best_idx = idx
model_fp_new = "{}{}-{}.pth".format(self.model_dir, epoch, lr)
torch.save(self.model.state_dict(), model_fp_new)
best_model = self.model
self.learning_rate = lrs[best_idx] * 0.8
self.optimizer = optimizers[best_idx]
self.model = best_model
test_loss = best_test
train_loss = best_train
return self.learning_rate, train_loss, test_loss
