iicd/autodl-projects
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

complete maml and remove vis_compare_algo

Browse Source
This commit is contained in:
D-X-Y
2021-05-10 11:19:18 +08:00
parent 755c7c90cf
commit 147da98f94
2 changed files with 103 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
exps/LFNA/basic-maml.py
View File

@@ -40,7 +40,7 @@ class MAML:
self.network.parameters(), lr=meta_lr, amsgrad=True
)
self.meta_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
self.meta_optimizer,
milestones=[
int(epochs * 0.25),
int(epochs * 0.5),
@@ -50,8 +50,8 @@ class MAML:
)
self.inner_lr = inner_lr
self.inner_step = inner_step
self._best_info = dict(state_dict=None, score=None)
print("There are {:} weights.".format(w_container.numel()))
self._best_info = dict(state_dict=None, iepoch=None, score=None)
print("There are {:} weights.".format(self.network.get_w_container().numel()))
def adapt(self, dataset):
# create a container for the future timestamp
@@ -61,7 +61,6 @@ class MAML:
y_hat = self.network.forward_with_container(dataset.x, container)
loss = self.criterion(y_hat, dataset.y)
grads = torch.autograd.grad(loss, container.parameters())
container = container.additive([-self.inner_lr * grad for grad in grads])
return container
@@ -73,23 +72,34 @@ class MAML:
return y_hat
def step(self):
torch.nn.utils.clip_grad_norm_(self.container.parameters(), 1.0)
torch.nn.utils.clip_grad_norm_(self.network.parameters(), 1.0)
self.meta_optimizer.step()
self.meta_lr_scheduler.step()
def zero_grad(self):
self.meta_optimizer.zero_grad()
def save_best(self, network, score):
def load_state_dict(self, state_dict):
self.criterion.load_state_dict(state_dict["criterion"])
self.network.load_state_dict(state_dict["network"])
self.meta_optimizer.load_state_dict(state_dict["meta_optimizer"])
self.meta_lr_scheduler.load_state_dict(state_dict["meta_lr_scheduler"])
def save_best(self, iepoch, score):
if self._best_info["score"] is None or self._best_info["score"] < score:
state_dict = dict(
criterion=criterion,
network=network.state_dict(),
criterion=self.criterion.state_dict(),
network=self.network.state_dict(),
meta_optimizer=self.meta_optimizer.state_dict(),
meta_lr_scheduler=self.meta_lr_scheduler.state_dict(),
)
self._best_info["state_dict"] = state_dict
self._best_info["score"] = score
self._best_info["iepoch"] = iepoch
is_best = True
else:
is_best = False
return self._best_info, is_best
def main(args):
@@ -111,8 +121,9 @@ def main(args):
# meta-training
per_epoch_time, start_time = AverageMeter(), time.time()
for iepoch in range(args.epochs):
# for iepoch in range(args.epochs):
iepoch = 0
while iepoch < args.epochs:
need_time = "Time Left: {:}".format(
convert_secs2time(per_epoch_time.avg * (args.epochs - iepoch), True)
)
@@ -122,9 +133,10 @@ def main(args):
)
maml.zero_grad()
meta_losses = []
batch_indexes, meta_losses = [], []
for ibatch in range(args.meta_batch):
sampled_timestamp = random.randint(0, train_time_bar)
batch_indexes.append("{:5d}".format(sampled_timestamp))
past_dataset = TimeData(
sampled_timestamp,
env_info["{:}-x".format(sampled_timestamp)],
@@ -135,7 +147,7 @@ def main(args):
env_info["{:}-x".format(sampled_timestamp + 1)],
env_info["{:}-y".format(sampled_timestamp + 1)],
)
future_container = maml.adapt(model, past_dataset)
future_container = maml.adapt(past_dataset)
future_y_hat = maml.predict(future_dataset.x, future_container)
future_loss = maml.criterion(future_y_hat, future_dataset.y)
meta_losses.append(future_loss)
@@ -143,14 +155,53 @@ def main(args):
meta_loss.backward()
maml.step()
logger.log("meta-loss: {:.4f}".format(meta_loss.item()))
logger.log(
"meta-loss: {:.4f} batch: {:}".format(
meta_loss.item(), ",".join(batch_indexes)
)
)
best_info, is_best = maml.save_best(iepoch, -meta_loss.item())
if is_best:
save_checkpoint(best_info, logger.path("best"), logger)
logger.log("Save the best into {:}".format(logger.path("best")))
if iepoch >= 10 and (
torch.isnan(meta_loss).item() or meta_loss.item() >= args.fail_thresh
):
xdata = torch.load(logger.path("best"))
maml.load_state_dict(xdata["state_dict"])
iepoch = xdata["iepoch"]
logger.log(
"The training failed, re-use the previous best epoch [{:}]".format(
iepoch
)
)
else:
iepoch = iepoch + 1
per_epoch_time.update(time.time() - start_time)
start_time = time.time()
import pdb
pdb.set_trace()
w_container_per_epoch = dict()
for idx in range(1, env_info["total"]):
past_dataset = TimeData(
idx - 1,
env_info["{:}-x".format(idx - 1)],
env_info["{:}-y".format(idx - 1)],
)
current_container = maml.adapt(past_dataset)
w_container_per_epoch[idx] = current_container.no_grad_clone()
with torch.no_grad():
current_x = env_info["{:}-x".format(idx)]
current_y = env_info["{:}-y".format(idx)]
current_y_hat = maml.predict(current_x, w_container_per_epoch[idx])
current_loss = maml.criterion(current_y_hat, current_y)
logger.log(
"meta-test: [{:03d}] -> loss={:.4f}".format(idx, current_loss.item())
)
save_checkpoint(
{"w_container_per_epoch": w_container_per_epoch},
logger.path(None) / "final-ckp.pth",
logger,
)
logger.log("-" * 200 + "\n")
logger.close()
@@ -179,9 +230,15 @@ if __name__ == "__main__":
parser.add_argument(
"--meta_lr",
type=float,
default=0.1,
default=0.05,
help="The learning rate for the MAML optimizer (default is Adam)",
)
parser.add_argument(
"--fail_thresh",
type=float,
default=1000,
help="The threshold for the failure, which we reuse the previous best model",
)
parser.add_argument(
"--inner_lr",
type=float,