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_epochs < num_epochs_cv:
self.num_epochs_cv = self.num_epochs
else:
self.num_epochs_cv = num_epochs_cv
if self.settings.BATCH_NORM: #i.e. generally larger learning_rate with BN
lrs = [0.006, 0.003, 0.015, 0.0015]
else:
lrs = [0.003, 0.001, 0.006, 0.0006]
res = []
optimizers = []
for idx, lr in enumerate(lrs):
print("lr:", lr)
ML_utils.set_seeds() #set seeds before init model
self.model = self.settings.AE_MODEL_TYPE(
**self.settings.get_kwargs())
self.optimizer = optim.Adam(self.model.parameters(), lr)
test_losses = []
train_losses = []
for epoch in range(self.num_epochs_cv):
train, test = self.train_one_epoch(epoch, 1, test_every,
self.num_epochs_cv)
if test:
test_losses.append(test)
train_losses.append(train)
df = pd.DataFrame(train_losses,
columns=[
"epoch", "reconstruction_err", "DA_MAE",
"DA_ratio_improve_MAE"
])
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]
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
def __init__(self, AE_settings, expdir, calc_DA_MAE=False, batch_sz=BATCH):
"""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
self.expdir = self.__init_expdir(expdir)
self.test_fp = self.expdir + "test.csv"
self.train_fp = self.expdir + "train.csv"
self.settings_fp = self.expdir + "settings.txt"
self.calc_DA_MAE = calc_DA_MAE
self.batch_sz = batch_sz
self.settings.batch_sz = batch_sz
ML_utils.set_seeds() #set seeds before init model
self.model = AE_settings.AE_MODEL_TYPE(**AE_settings.get_kwargs())
print("Number of parameters:",
sum(p.numel() for p in self.model.parameters()))
self.device = ML_utils.get_device()
def main():
dir = "/data/home/jfm1118/DA/experiments/train_DA_Pressure/2-l4NBN/" # 299.pth"
model, settings = ML_utils.load_model_and_settings_from_dir(dir)
settings.TOL = 1e-3
settings.SAVE = False
settings.export_env_vars()
#Load data
loader, splitter = GetData(), SplitData()
X = loader.get_X(settings)
train_X, test_X, u_c, X, mean, std = splitter.train_test_DA_split_maybe_normalize(X, settings)
L1 = torch.nn.L1Loss(reduction='sum')
L2 = torch.nn.MSELoss(reduction="sum")
datasets = {"train": train_X,
"test": test_X,
"u_c": u_c,
"mean": mean,
"u_0": np.zeros_like(u_c)}
# datasets = {
# "train1": train_X[:1],
# "train2": train_X[:2],
# "train3": train_X[:3],
# "train4": train_X[:4],
# "train5": train_X[:5],
# "train6": train_X[:6],
# }
for name, data in datasets.items():
if len(data.shape) in [1, 3]:
num_states = 1
else:
num_states = data.shape[0]
device = ML_utils.get_device()
data_tensor = torch.Tensor(data)
data_tensor = data_tensor.to(device)
DA_pipeline = DAPipeline(settings, model)
encoder = DA_pipeline.data.get("encoder")
decoder = DA_pipeline.data.get("decoder")
data_hat = decoder(encoder(data))
data_hat = torch.Tensor(data_hat)
data_hat = data_hat.to(device)
l1 = L1(data_hat, data_tensor)
l2 = L2(data_hat, data_tensor)
print("name: {}, \nL1: {:.2f}\nL2: {:.2f}\n".format(name, l1 / num_states, l2 / num_states,))
print()
def main():
model, settings = ML_utils.load_model_and_settings_from_dir(
DIR, device_idx=GPU_DEVICE)
device = ML_utils.get_device(device_idx=GPU_DEVICE)
initializer = VDAInit(settings, model)
vdadata = initializer.run()
encoder = vdadata.get("encoder")
loader, splitter = GetData(), SplitData()
X = loader.get_X(settings)
train_X, test_X, u_c, X, mean, std = splitter.train_test_DA_split_maybe_normalize(
X, settings)
settings.AUGMENTATION = False
train_loader, test_loader = loader.get_train_test_loaders(settings,
32,
num_workers=6)
Z = []
for Bidx, x in enumerate(test_loader):
x, = x
z = encoder(x)
Z.append(z)
Z_test = np.concatenate(Z, axis=0)
print(Z_test.shape)
#Compare all test data with all other test data
tot, tot_abs = 0, 0
number = 0
for idx1 in range(Z_test.shape[0]):
for idx2 in range(Z_test.shape[0]):
if idx2 <= idx1:
continue
res = compare_overlap(Z_test[idx1], Z_test[idx2])
tot_abs += np.abs(res)
tot += res
number += 1
if number % 10000 == 0 and number != 0:
print(number, tot_abs / number, tot / number)
print("Final", tot_abs / number, tot / number)
def select_obs(self, mode, vec, frac=None):
"""Selects and return a subset of observations and their indexes
from vec according to a user selected mode"""
npoints = self.get_npoints_from_shape(vec.shape)
if mode == "rand":
# Define observations as a random subset of the control state.
nobs = int(frac * npoints) #number of observations
ML_utils.set_seeds(
seed=self.settings.SEED
) #set seeds so that the selected subset is the same every time
obs_idx = random.sample(
range(npoints), nobs) #select nobs integers w/o replacement
observations = np.take(vec, obs_idx)
elif mode == "single_max":
nobs = 1
obs_idx = np.argmax(vec)
obs_idx = [obs_idx]
observations = np.take(vec, obs_idx)
return observations, obs_idx, nobs
def training_loop_AE(self,
start_epoch,
num_epoch,
device=None,
print_every=2,
test_every=5,
save_every=5,
model_dir=None):
"""Runs a torch AE model training loop.
NOTE: Ensure that the loss_fn is in mode "sum"
"""
model = self.model
self.model_dir = model_dir
if device == None:
device = ML_utils.get_device()
self.device = device
ML_utils.set_seeds()
train_losses = []
test_losses = []
epoch = num_epoch - 1 #for case where no training occurs
for epoch in range(start_epoch, num_epoch):
train_loss, test_loss = self.train_one_epoch(
epoch, print_every, test_every, num_epoch)
train_losses.append(train_loss)
if test_loss:
test_losses.append(test_loss)
if epoch % save_every != 0 and self.model_dir != None:
#Save model (if new model hasn't just been saved)
model_fp_new = "{}{}.pth".format(self.model_dir, epoch)
torch.save(model.state_dict(), model_fp_new)
return train_losses, test_losses
def vda_setup(self, settings):
"""Generates matrices for VarDA. All returned matrices are in the
(n X M) format (as typical in VarDA) although when
settings.THREE_DIM = True, some are 4-dimensional"""
data = {}
loader = GetData()
splitter = SplitData()
X = loader.get_X(settings)
train_X, test_X, u_c, X, mean, std = splitter.train_test_DA_split_maybe_normalize(
X, settings)
V = self.create_V_from_X(train_X, settings)
#Deal with dimensions:
#currently dim are: (M x n ) or (M x nx x ny x nz)
#change to (n x M) or (nx x ny x nz x M)
if not settings.THREE_DIM:
X = X.T
train_X = train_X.T
test_X = test_X.T
V = V.T
else:
# can flatten V as it will only be used for SVD (which ignores 3D location)
V = V.reshape((V.shape[0], -1)).T
X = np.moveaxis(X, 0, 3)
train_X = np.moveaxis(train_X, 0, 3)
test_X = np.moveaxis(test_X, 0, 3)
# We will take initial condition u_0, as mean of historical data
if settings.NORMALIZE:
u_0 = np.zeros(settings.get_n()) #since the data is mean centred
else:
u_0 = mean
#flatten 3D vectors:
u_c = u_c.flatten()
std = std.flatten()
mean = mean.flatten()
u_0 = u_0.flatten()
observations, obs_idx, nobs = self.select_obs(
settings.OBS_MODE, u_c,
settings.OBS_FRAC) #options are specific for rand
#Now define quantities required for 3D-VarDA - see Algorithm 1 in Rossella et al (2019)
H_0 = self.create_H(obs_idx, settings.get_n(), nobs,
settings.THREE_DIM)
d = observations - H_0 @ u_0 #'d' in literature
#R_inv = self.create_R_inv(OBS_VARIANCE, nobs)
device = ML_utils.get_device(False)
data = {
"d": d,
"G": H_0,
"V": V,
"observations": observations,
"u_c": u_c,
"u_0": u_0,
"X": X,
"train_X": train_X,
"test_X": test_X,
"std": std,
"mean": mean,
"device": device
}
return data, std, mean
def train(self,
num_epoch=100,
learning_rate=0.00025,
print_every=2,
test_every=5):
self.learning_rate = learning_rate
self.num_epochs = num_epoch
settings = self.settings
if settings.SAVE == True:
self.model_dir = self.expdir
else:
self.model_dir = None
#data
loader = GetData()
splitter = SplitData()
X = loader.get_X(settings)
self.train_X, self.test_X, DA_u_c, X_norm, mean, std = splitter.train_test_DA_split_maybe_normalize(
X, settings)
#Add Channel if we are in 3D case
if settings.THREE_DIM:
self.train_X = np.expand_dims(self.train_X, 1)
self.test_X = np.expand_dims(self.test_X, 1)
#Dataloaders
train_dataset = TensorDataset(torch.Tensor(self.train_X))
self.train_loader = DataLoader(train_dataset,
self.batch_sz,
shuffle=True,
num_workers=6)
test_dataset = TensorDataset(torch.Tensor(self.test_X))
test_batch_sz = min(self.test_X.shape[0], self.batch_sz)
self.test_loader = DataLoader(test_dataset, test_batch_sz)
device = ML_utils.get_device()
self.loss_fn = torch.nn.L1Loss(reduction='sum')
self.learning_rate, train_losses, test_losses = self.__maybe_cross_val_lr(
test_every=test_every)
settings.learning_rate = self.learning_rate #for logging
train_losses_, test_losses_ = self.training_loop_AE(
self.num_epochs_cv,
self.num_epochs,
device,
print_every=print_every,
test_every=test_every,
model_dir=self.model_dir)
if train_losses_:
train_losses.extend(train_losses_)
if test_losses_:
test_losses.extend(test_losses_)
#Save results and settings file (so that it can be exactly reproduced)
if settings.SAVE == True:
self.to_csv(train_losses, self.train_fp)
self.to_csv(test_losses, self.test_fp)
with open(self.settings_fp, "wb") as f:
pickle.dump(settings, f)
return self.model