def main(args):
dataset_name = args.dataset
model_name = args.model
n_inner_iter = args.adaptation_steps
batch_size = args.batch_size
save_model_file = args.save_model_file
load_model_file = args.load_model_file
lower_trial = args.lower_trial
upper_trial = args.upper_trial
is_test = args.is_test
stopping_patience = args.stopping_patience
epochs = args.epochs
fast_lr = args.learning_rate
slow_lr = args.meta_learning_rate
noise_level = args.noise_level
noise_type = args.noise_type
resume = args.resume
first_order = False
inner_loop_grad_clip = 20
task_size = 50
output_dim = 1
checkpoint_freq = 10
horizon = 10
##test
meta_info = {
"POLLUTION": [5, 50, 14],
"HR": [32, 50, 13],
"BATTERY": [20, 50, 3]
}
assert model_name in ("FCN", "LSTM"), "Model was not correctly specified"
assert dataset_name in ("POLLUTION", "HR", "BATTERY")
window_size, task_size, input_dim = meta_info[dataset_name]
grid = [0., noise_level]
output_directory = "output/"
train_data_ML = pickle.load(
open(
"../../Data/TRAIN-" + dataset_name + "-W" + str(window_size) +
"-T" + str(task_size) + "-ML.pickle", "rb"))
validation_data_ML = pickle.load(
open(
"../../Data/VAL-" + dataset_name + "-W" + str(window_size) + "-T" +
str(task_size) + "-ML.pickle", "rb"))
test_data_ML = pickle.load(
open(
"../../Data/TEST-" + dataset_name + "-W" + str(window_size) +
"-T" + str(task_size) + "-ML.pickle", "rb"))
for trial in range(lower_trial, upper_trial):
output_directory = "../../Models/" + dataset_name + "_" + model_name + "_MAML/" + str(
trial) + "/"
save_model_file_ = output_directory + save_model_file
save_model_file_encoder = output_directory + "encoder_" + save_model_file
load_model_file_ = output_directory + load_model_file
checkpoint_file = output_directory + "checkpoint_" + save_model_file.split(
".")[0]
try:
os.mkdir(output_directory)
except OSError as error:
print(error)
with open(output_directory + "/results2.txt", "a+") as f:
f.write("Learning rate :%f \n" % fast_lr)
f.write("Meta-learning rate: %f \n" % slow_lr)
f.write("Adaptation steps: %f \n" % n_inner_iter)
f.write("Noise level: %f \n" % noise_level)
if model_name == "LSTM":
model = LSTMModel(batch_size=batch_size,
seq_len=window_size,
input_dim=input_dim,
n_layers=2,
hidden_dim=120,
output_dim=output_dim)
model2 = LinearModel(120, 1)
optimizer = torch.optim.Adam(list(model.parameters()) +
list(model2.parameters()),
lr=slow_lr)
loss_func = mae
#loss_func = nn.SmoothL1Loss()
#loss_func = nn.MSELoss()
initial_epoch = 0
#torch.backends.cudnn.enabled = False
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
meta_learner = MetaLearner(model2, optimizer, fast_lr, loss_func,
first_order, n_inner_iter,
inner_loop_grad_clip, device)
model.to(device)
early_stopping = EarlyStopping(patience=stopping_patience,
model_file=save_model_file_encoder,
verbose=True)
early_stopping2 = EarlyStopping(patience=stopping_patience,
model_file=save_model_file_,
verbose=True)
if resume:
checkpoint = torch.load(checkpoint_file)
model.load_state_dict(checkpoint["model"])
meta_learner.load_state_dict(checkpoint["meta_learner"])
initial_epoch = checkpoint["epoch"]
best_score = checkpoint["best_score"]
counter = checkpoint["counter_stopping"]
early_stopping.best_score = best_score
early_stopping2.best_score = best_score
early_stopping.counter = counter
early_stopping2.counter = counter
total_tasks, task_size, window_size, input_dim = train_data_ML.x.shape
accum_mean = 0.0
for epoch in range(initial_epoch, epochs):
model.zero_grad()
meta_learner._model.zero_grad()
#train
batch_idx = np.random.randint(0, total_tasks - 1, batch_size)
#for batch_idx in range(0, total_tasks-1, batch_size):
x_spt, y_spt = train_data_ML[batch_idx]
x_qry, y_qry = train_data_ML[batch_idx + 1]
x_spt, y_spt = to_torch(x_spt), to_torch(y_spt)
x_qry = to_torch(x_qry)
y_qry = to_torch(y_qry)
# data augmentation
epsilon = grid[np.random.randint(0, len(grid))]
if noise_type == "additive":
y_spt = y_spt + epsilon
y_qry = y_qry + epsilon
else:
y_spt = y_spt * (1 + epsilon)
y_qry = y_qry * (1 + epsilon)
train_tasks = [
Task(model.encoder(x_spt[i]), y_spt[i])
for i in range(x_spt.shape[0])
]
val_tasks = [
Task(model.encoder(x_qry[i]), y_qry[i])
for i in range(x_qry.shape[0])
]
adapted_params = meta_learner.adapt(train_tasks)
mean_loss = meta_learner.step(adapted_params,
val_tasks,
is_training=True)
#accum_mean += mean_loss.cpu().detach().numpy()
#progressBar(batch_idx, total_tasks, 100)
#print(accum_mean/(batch_idx+1))
#test
val_error = test(validation_data_ML, meta_learner, model, device,
noise_level)
test_error = test(test_data_ML, meta_learner, model, device, 0.0)
print("Epoch:", epoch)
print("Val error:", val_error)
print("Test error:", test_error)
early_stopping(val_error, model)
early_stopping2(val_error, meta_learner)
#checkpointing
if epochs % checkpoint_freq == 0:
torch.save(
{
"epoch": epoch,
"model": model.state_dict(),
"meta_learner": meta_learner.state_dict(),
"best_score": early_stopping2.best_score,
"counter_stopping": early_stopping2.counter
}, checkpoint_file)
if early_stopping.early_stop:
print("Early stopping")
break
print("hallo")
model.load_state_dict(torch.load(save_model_file_encoder))
model2.load_state_dict(
torch.load(save_model_file_)["model_state_dict"])
meta_learner = MetaLearner(model2, optimizer, fast_lr, loss_func,
first_order, n_inner_iter,
inner_loop_grad_clip, device)
validation_error = test(validation_data_ML,
meta_learner,
model,
device,
noise_level=0.0)
test_error = test(test_data_ML,
meta_learner,
model,
device,
noise_level=0.0)
validation_error_h1 = test(validation_data_ML,
meta_learner,
model,
device,
noise_level=0.0,
horizon=1)
test_error_h1 = test(test_data_ML,
meta_learner,
model,
device,
noise_level=0.0,
horizon=1)
model.load_state_dict(torch.load(save_model_file_encoder))
model2.load_state_dict(
torch.load(save_model_file_)["model_state_dict"])
meta_learner2 = MetaLearner(model2, optimizer, fast_lr, loss_func,
first_order, 0, inner_loop_grad_clip,
device)
validation_error_h0 = test(validation_data_ML,
meta_learner2,
model,
device,
noise_level=0.0,
horizon=1)
test_error_h0 = test(test_data_ML,
meta_learner2,
model,
device,
noise_level=0.0,
horizon=1)
model.load_state_dict(torch.load(save_model_file_encoder))
model2.load_state_dict(
torch.load(save_model_file_)["model_state_dict"])
meta_learner2 = MetaLearner(model2, optimizer, fast_lr, loss_func,
first_order, n_inner_iter,
inner_loop_grad_clip, device)
validation_error_mae = test(validation_data_ML, meta_learner2, model,
device, 0.0)
test_error_mae = test(test_data_ML, meta_learner2, model, device, 0.0)
print("test_error_mae", test_error_mae)
with open(output_directory + "/results2.txt", "a+") as f:
f.write("Test error: %f \n" % test_error)
f.write("Validation error: %f \n" % validation_error)
f.write("Test error h1: %f \n" % test_error_h1)
f.write("Validation error h1: %f \n" % validation_error_h1)
f.write("Test error h0: %f \n" % test_error_h0)
f.write("Validation error h0: %f \n" % validation_error_h0)
f.write("Test error mae: %f \n" % test_error_mae)
f.write("Validation error mae: %f \n" % validation_error_mae)
print(test_error)
print(validation_error)
def test(maml,
model_name,
dataset_name,
test_data_ML,
adaptation_steps,
learning_rate,
with_early_stopping=False,
horizon=10):
total_tasks_test = len(test_data_ML)
error_list = []
learner = maml.clone() # Creates a clone of model
accum_error = 0.0
count = 0
task_size = test_data_ML.x.shape[-3]
input_dim = test_data_ML.x.shape[-1]
window_size = test_data_ML.x.shape[-2]
output_dim = test_data_ML.y.shape[-1]
for task in range(0, (total_tasks_test - horizon - 1),
total_tasks_test // 100):
temp_file_idx = test_data_ML.file_idx[task:task + horizon + 1]
if (len(np.unique(temp_file_idx)) > 1):
continue
if model_name == "LSTM":
model2 = LSTMModel(batch_size=None,
seq_len=None,
input_dim=input_dim,
n_layers=2,
hidden_dim=120,
output_dim=1)
elif model_name == "FCN":
kernels = [8, 5, 3] if window_size != 5 else [4, 2, 1]
model2 = FCN(time_steps=window_size,
channels=[input_dim, 128, 128, 128],
kernels=kernels)
model2.cuda()
model2.load_state_dict(copy.deepcopy(maml.module.state_dict()))
opt2 = optim.SGD(model2.parameters(), lr=learning_rate)
if with_early_stopping:
x_spt, y_spt = test_data_ML[task]
x_spt_val = x_spt[int(task_size * 0.8):]
y_spt_val = y_spt[int(task_size * 0.8):]
x_spt = x_spt[:int(task_size * 0.8)]
y_spt = y_spt[:int(task_size * 0.8)]
x_qry = test_data_ML.x[(task + 1):(task + 1 + horizon)].reshape(
-1, window_size, input_dim)
y_qry = test_data_ML.y[(task + 1):(task + 1 + horizon)].reshape(
-1, output_dim)
if model_name == "FCN":
x_qry = np.transpose(x_qry, [0, 2, 1])
x_spt = np.transpose(x_spt, [0, 2, 1])
x_spt_val = np.transpose(x_spt_val, [0, 2, 1])
x_spt, y_spt = to_torch(x_spt), to_torch(y_spt)
x_spt_val, y_spt_val = to_torch(x_spt_val), to_torch(y_spt_val)
x_qry = to_torch(x_qry)
y_qry = to_torch(y_qry)
else:
x_spt, y_spt = test_data_ML[task]
x_qry = test_data_ML.x[(task + 1):(task + 1 + horizon)].reshape(
-1, window_size, input_dim)
y_qry = test_data_ML.y[(task + 1):(task + 1 + horizon)].reshape(
-1, output_dim)
if model_name == "FCN":
x_qry = np.transpose(x_qry, [0, 2, 1])
x_spt = np.transpose(x_spt, [0, 2, 1])
x_spt, y_spt = to_torch(x_spt), to_torch(y_spt)
x_qry = to_torch(x_qry)
y_qry = to_torch(y_qry)
early_stopping = EarlyStopping(patience=2,
model_file="temp/temp_file_" +
model_name + ".pt",
verbose=True)
#model2.eval()
for step in range(adaptation_steps):
model2.zero_grad()
opt2.zero_grad()
#model2.train()
pred = model2(x_spt)
error = mae(pred, y_spt)
error.backward()
opt2.step()
if with_early_stopping:
with torch.no_grad():
#model2.eval()
pred = model2(x_spt_val)
error = mae(pred, y_spt_val)
early_stopping(error, model2)
if early_stopping.early_stop:
print("Early stopping")
break
if with_early_stopping:
model2.load_state_dict(
torch.load("temp/temp_file_" + model_name + ".pt"))
#model2.eval()
pred = model2(x_qry)
error = mae(pred, y_qry)
accum_error += error.data
count += 1
error = accum_error / count
return error