iicd/xautodl
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update ablation for GeMOSA

Browse Source
This commit is contained in:
D-X-Y 2021-05-27 19:47:08 +08:00
parent 08337138f1
commit 5dd75696c9
4 changed files with 304 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
exps/GeMOSA/basic-same.py
View File

@ -4,6 +4,7 @@
# python exps/GeMOSA/basic-same.py --env_version v1 --hidden_dim 16 --epochs 500 --init_lr 0.1 --device cuda
# python exps/GeMOSA/basic-same.py --env_version v2 --hidden_dim 16 --epochs 500 --init_lr 0.1 --device cuda
# python exps/GeMOSA/basic-same.py --env_version v3 --hidden_dim 32 --epochs 1000 --init_lr 0.05 --device cuda
# python exps/GeMOSA/basic-same.py --env_version v4 --hidden_dim 32 --epochs 1000 --init_lr 0.05 --device cuda
#####################################################
import sys, time, copy, torch, random, argparse
from tqdm import tqdm
@ -28,7 +29,12 @@ from xautodl.log_utils import AverageMeter, convert_secs2time
from xautodl.utils import split_str2indexes
from xautodl.procedures.advanced_main import basic_train_fn, basic_eval_fn
from xautodl.procedures.metric_utils import SaveMetric, MSEMetric, ComposeMetric
from xautodl.procedures.metric_utils import (
SaveMetric,
MSEMetric,
Top1AccMetric,
ComposeMetric,
)
from xautodl.datasets.synthetic_core import get_synthetic_env
from xautodl.models.xcore import get_model
@ -57,6 +63,17 @@ def main(args):
logger.log("The total enviornment: {:}".format(env))
w_containers = dict()
if env.meta_info["task"] == "regression":
criterion = torch.nn.MSELoss()
metric_cls = MSEMetric
elif env.meta_info["task"] == "classification":
criterion = torch.nn.CrossEntropyLoss()
metric_cls = Top1AccMetric
else:
raise ValueError(
"This task ({:}) is not supported.".format(all_env.meta_info["task"])
)
per_timestamp_time, start_time = AverageMeter(), time.time()
for idx, (future_time, (future_x, future_y)) in enumerate(env):
@ -79,7 +96,6 @@ def main(args):
print(model)
# build optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.init_lr, amsgrad=True)
criterion = torch.nn.MSELoss()
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[
@ -89,7 +105,7 @@ def main(args):
],
gamma=0.3,
)
train_metric = MSEMetric()
train_metric = metric_cls(True)
best_loss, best_param = None, None
for _iepoch in range(args.epochs):
preds = model(historical_x)
@ -108,19 +124,19 @@ def main(args):
train_metric(preds, historical_y)
train_results = train_metric.get_info()
metric = ComposeMetric(MSEMetric(), SaveMetric())
xmetric = ComposeMetric(metric_cls(True), SaveMetric())
eval_dataset = torch.utils.data.TensorDataset(
future_x.to(args.device), future_y.to(args.device)
)
eval_loader = torch.utils.data.DataLoader(
eval_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0
)
results = basic_eval_fn(eval_loader, model, metric, logger)
results = basic_eval_fn(eval_loader, model, xmetric, logger)
log_str = (
"[{:}]".format(time_string())
+ " [{:04d}/{:04d}]".format(idx, len(env))
+ " train-mse: {:.5f}, eval-mse: {:.5f}".format(
train_results["mse"], results["mse"]
+ " train-score: {:.5f}, eval-score: {:.5f}".format(
train_results["score"], results["score"]
)
)
logger.log(log_str)

27
exps/GeMOSA/main.py
View File

@ -1,12 +1,16 @@
#####################################################
# Learning to Generate Model One Step Ahead #
#####################################################
##########################################################
# Learning to Efficiently Generate Models One Step Ahead #
##########################################################
# <----> run on CPU
# python exps/GeMOSA/main.py --env_version v1 --workers 0
# <----> run on a GPU
# python exps/GeMOSA/main.py --env_version v1 --lr 0.002 --hidden_dim 16 --meta_batch 256 --device cuda
# python exps/GeMOSA/main.py --env_version v2 --lr 0.002 --hidden_dim 16 --meta_batch 256 --device cuda
# python exps/GeMOSA/main.py --env_version v3 --lr 0.002 --hidden_dim 32 --time_dim 32 --meta_batch 256 --device cuda
# python exps/GeMOSA/main.py --env_version v4 --lr 0.002 --hidden_dim 32 --time_dim 32 --meta_batch 256 --device cuda
#####################################################
# <----> ablation commands
# python exps/GeMOSA/main.py --env_version v4 --lr 0.002 --hidden_dim 32 --time_dim 32 --meta_batch 256 --ablation old --device cuda
##########################################################
import sys, time, copy, torch, random, argparse
from tqdm import tqdm
from copy import deepcopy
@ -36,6 +40,7 @@ from xautodl.models.xcore import get_model
from xautodl.procedures.metric_utils import MSEMetric, Top1AccMetric
from meta_model import MetaModelV1
from meta_model_ablation import MetaModel_TraditionalAtt
def online_evaluate(
@ -230,7 +235,13 @@ def main(args):
# pre-train the hypernetwork
timestamps = trainval_env.get_timestamp(None)
meta_model = MetaModelV1(
if args.ablation is None:
MetaModel_cls = MetaModelV1
elif args.ablation == "old":
MetaModel_cls = MetaModel_TraditionalAtt
else:
raise ValueError("Unknown ablation : {:}".format(args.ablation))
meta_model = MetaModel_cls(
shape_container,
args.layer_dim,
args.time_dim,
@ -373,6 +384,9 @@ if __name__ == "__main__":
parser.add_argument(
"--workers", type=int, default=4, help="The number of workers in parallel."
)
parser.add_argument(
"--ablation", type=str, default=None, help="The ablation indicator."
)
parser.add_argument(
"--device",
type=str,
@ -385,7 +399,7 @@ if __name__ == "__main__":
if args.rand_seed is None or args.rand_seed < 0:
args.rand_seed = random.randint(1, 100000)
assert args.save_dir is not None, "The save dir argument can not be None"
args.save_dir = "{:}-bs{:}-d{:}_{:}_{:}-s{:}-lr{:}-wd{:}-e{:}-env{:}".format(
args.save_dir = "{:}-bs{:}-d{:}_{:}_{:}-s{:}-lr{:}-wd{:}-e{:}-ab{:}-env{:}".format(
args.save_dir,
args.meta_batch,
args.hidden_dim,
@ -395,6 +409,7 @@ if __name__ == "__main__":
args.lr,
args.weight_decay,
args.epochs,
args.ablation,
args.env_version,
)
main(args)

3
exps/GeMOSA/meta_model.py
View File

@ -1,6 +1,3 @@
#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.04 #
#####################################################
import torch
import torch.nn.functional as F

260
exps/GeMOSA/meta_model_ablation.py Normal file
View File

@ -0,0 +1,260 @@
#
# This is used for the ablation studies:
# The meta-model in this file uses the traditional attention in
# transformer.
#
import torch
import torch.nn.functional as F
from xautodl.xlayers import super_core
from xautodl.xlayers import trunc_normal_
from xautodl.models.xcore import get_model
class MetaModel_TraditionalAtt(super_core.SuperModule):
"""Learning to Generate Models One Step Ahead (Meta Model Design)."""
def __init__(
self,
shape_container,
layer_dim,
time_dim,
meta_timestamps,
dropout: float = 0.1,
seq_length: int = None,
interval: float = None,
thresh: float = None,
):
super(MetaModel_TraditionalAtt, self).__init__()
self._shape_container = shape_container
self._num_layers = len(shape_container)
self._numel_per_layer = []
for ilayer in range(self._num_layers):
self._numel_per_layer.append(shape_container[ilayer].numel())
self._raw_meta_timestamps = meta_timestamps
assert interval is not None
self._interval = interval
self._thresh = interval * seq_length if thresh is None else thresh
self.register_parameter(
"_super_layer_embed",
torch.nn.Parameter(torch.Tensor(self._num_layers, layer_dim)),
)
self.register_parameter(
"_super_meta_embed",
torch.nn.Parameter(torch.Tensor(len(meta_timestamps), time_dim)),
)
self.register_buffer("_meta_timestamps", torch.Tensor(meta_timestamps))
self._time_embed_dim = time_dim
self._append_meta_embed = dict(fixed=None, learnt=None)
self._append_meta_timestamps = dict(fixed=None, learnt=None)
self._tscalar_embed = super_core.SuperDynamicPositionE(
time_dim, scale=1 / interval
)
# build transformer
self._trans_att = super_core.SuperQKVAttention(
in_q_dim=time_dim,
in_k_dim=time_dim,
in_v_dim=time_dim,
num_heads=4,
proj_dim=time_dim,
qkv_bias=True,
attn_drop=None,
proj_drop=dropout,
)
model_kwargs = dict(
config=dict(model_type="dual_norm_mlp"),
input_dim=layer_dim + time_dim,
output_dim=max(self._numel_per_layer),
hidden_dims=[(layer_dim + time_dim) * 2] * 3,
act_cls="gelu",
norm_cls="layer_norm_1d",
dropout=dropout,
)
self._generator = get_model(**model_kwargs)
# initialization
trunc_normal_(
[self._super_layer_embed, self._super_meta_embed],
std=0.02,
)
def get_parameters(self, time_embed, attention, generator):
parameters = []
if time_embed:
parameters.append(self._super_meta_embed)
if attention:
parameters.extend(list(self._trans_att.parameters()))
if generator:
parameters.append(self._super_layer_embed)
parameters.extend(list(self._generator.parameters()))
return parameters
@property
def meta_timestamps(self):
with torch.no_grad():
meta_timestamps = [self._meta_timestamps]
for key in ("fixed", "learnt"):
if self._append_meta_timestamps[key] is not None:
meta_timestamps.append(self._append_meta_timestamps[key])
return torch.cat(meta_timestamps)
@property
def super_meta_embed(self):
meta_embed = [self._super_meta_embed]
for key in ("fixed", "learnt"):
if self._append_meta_embed[key] is not None:
meta_embed.append(self._append_meta_embed[key])
return torch.cat(meta_embed)
def create_meta_embed(self):
param = torch.Tensor(1, self._time_embed_dim)
trunc_normal_(param, std=0.02)
param = param.to(self._super_meta_embed.device)
param = torch.nn.Parameter(param, True)
return param
def get_closest_meta_distance(self, timestamp):
with torch.no_grad():
distances = torch.abs(self.meta_timestamps - timestamp)
return torch.min(distances).item()
def replace_append_learnt(self, timestamp, meta_embed):
self._append_meta_timestamps["learnt"] = timestamp
self._append_meta_embed["learnt"] = meta_embed
@property
def meta_length(self):
return self.meta_timestamps.numel()
def clear_fixed(self):
self._append_meta_timestamps["fixed"] = None
self._append_meta_embed["fixed"] = None
def clear_learnt(self):
self.replace_append_learnt(None, None)
def append_fixed(self, timestamp, meta_embed):
with torch.no_grad():
device = self._super_meta_embed.device
timestamp = timestamp.detach().clone().to(device)
meta_embed = meta_embed.detach().clone().to(device)
if self._append_meta_timestamps["fixed"] is None:
self._append_meta_timestamps["fixed"] = timestamp
else:
self._append_meta_timestamps["fixed"] = torch.cat(
(self._append_meta_timestamps["fixed"], timestamp), dim=0
)
if self._append_meta_embed["fixed"] is None:
self._append_meta_embed["fixed"] = meta_embed
else:
self._append_meta_embed["fixed"] = torch.cat(
(self._append_meta_embed["fixed"], meta_embed), dim=0
)
def gen_time_embed(self, timestamps):
# timestamps is a batch of timestamps
[B] = timestamps.shape
# batch, seq = timestamps.shape
timestamps = timestamps.view(-1, 1)
meta_timestamps, meta_embeds = self.meta_timestamps, self.super_meta_embed
timestamp_v_embed = meta_embeds.unsqueeze(dim=0)
timestamp_q_embed = self._tscalar_embed(timestamps)
timestamp_k_embed = self._tscalar_embed(meta_timestamps.view(1, -1))
# create the mask
mask = (
torch.unsqueeze(timestamps, dim=-1) <= meta_timestamps.view(1, 1, -1)
) | (
torch.abs(
torch.unsqueeze(timestamps, dim=-1) - meta_timestamps.view(1, 1, -1)
)
> self._thresh
)
timestamp_embeds = self._trans_att(
timestamp_q_embed, timestamp_k_embed, timestamp_v_embed, mask
)
return timestamp_embeds[:, -1, :]
def gen_model(self, time_embeds):
B, _ = time_embeds.shape
# create joint embed
num_layer, _ = self._super_layer_embed.shape
# The shape of `joint_embed` is batch * num-layers * input-dim
joint_embeds = torch.cat(
(
time_embeds.view(B, 1, -1).expand(-1, num_layer, -1),
self._super_layer_embed.view(1, num_layer, -1).expand(B, -1, -1),
),
dim=-1,
)
batch_weights = self._generator(joint_embeds)
batch_containers = []
for weights in torch.split(batch_weights, 1):
batch_containers.append(
self._shape_container.translate(torch.split(weights.squeeze(0), 1))
)
return batch_containers
def forward_raw(self, timestamps, time_embeds, tembed_only=False):
raise NotImplementedError
def forward_candidate(self, input):
raise NotImplementedError
def easy_adapt(self, timestamp, time_embed):
with torch.no_grad():
timestamp = torch.Tensor([timestamp]).to(self._meta_timestamps.device)
self.replace_append_learnt(None, None)
self.append_fixed(timestamp, time_embed)
def adapt(self, base_model, criterion, timestamp, x, y, lr, epochs, init_info):
distance = self.get_closest_meta_distance(timestamp)
if distance + self._interval * 1e-2 <= self._interval:
return False, None
x, y = x.to(self._meta_timestamps.device), y.to(self._meta_timestamps.device)
with torch.set_grad_enabled(True):
new_param = self.create_meta_embed()
optimizer = torch.optim.Adam(
[new_param], lr=lr, weight_decay=1e-5, amsgrad=True
)
timestamp = torch.Tensor([timestamp]).to(new_param.device)
self.replace_append_learnt(timestamp, new_param)
self.train()
base_model.train()
if init_info is not None:
best_loss = init_info["loss"]
new_param.data.copy_(init_info["param"].data)
else:
best_loss = 1e9
with torch.no_grad():
best_new_param = new_param.detach().clone()
for iepoch in range(epochs):
optimizer.zero_grad()
time_embed = self.gen_time_embed(timestamp.view(1))
match_loss = criterion(new_param, time_embed)
[container] = self.gen_model(new_param.view(1, -1))
y_hat = base_model.forward_with_container(x, container)
meta_loss = criterion(y_hat, y)
loss = meta_loss + match_loss
loss.backward()
optimizer.step()
if meta_loss.item() < best_loss:
with torch.no_grad():
best_loss = meta_loss.item()
best_new_param = new_param.detach().clone()
self.easy_adapt(timestamp, best_new_param)
return True, best_loss
def extra_repr(self) -> str:
return "(_super_layer_embed): {:}, (_super_meta_embed): {:}, (_meta_timestamps): {:}".format(
list(self._super_layer_embed.shape),
list(self._super_meta_embed.shape),
list(self._meta_timestamps.shape),
)