Try a different model / LFNA
This commit is contained in:
D-X-Y
parent
25dc78a7ce
commit
9135667cc1
@ -99,7 +99,7 @@ def online_evaluate(env, meta_model, base_model, criterion, args, logger):
|
||||
future_time = timestamp.item()
|
||||
time_seqs = [
|
||||
future_time - iseq * env.timestamp_interval
|
||||
for iseq in range(args.seq_length * 2)
|
||||
for iseq in range(args.seq_length)
|
||||
]
|
||||
time_seqs.reverse()
|
||||
with torch.no_grad():
|
||||
@ -107,30 +107,26 @@ def online_evaluate(env, meta_model, base_model, criterion, args, logger):
|
||||
base_model.eval()
|
||||
time_seqs = torch.Tensor(time_seqs).view(1, -1).to(args.device)
|
||||
[seq_containers], _ = meta_model(time_seqs, None)
|
||||
# For Debug
|
||||
for idx in range(time_seqs.numel()):
|
||||
future_container = seq_containers[idx]
|
||||
_, (future_x, future_y) = env(time_seqs[0, idx].item())
|
||||
future_x, future_y = future_x.to(args.device), future_y.to(args.device)
|
||||
future_y_hat = base_model.forward_with_container(
|
||||
future_x, future_container
|
||||
)
|
||||
future_loss = criterion(future_y_hat, future_y)
|
||||
logger.log(
|
||||
"--> time={:.4f} -> loss={:.4f}".format(
|
||||
time_seqs[0, idx].item(), future_loss.item()
|
||||
)
|
||||
)
|
||||
future_container = seq_containers[-1]
|
||||
future_x, future_y = future_x.to(args.device), future_y.to(args.device)
|
||||
future_y_hat = base_model.forward_with_container(future_x, future_container)
|
||||
future_loss = criterion(future_y_hat, future_y)
|
||||
logger.log(
|
||||
"[ONLINE] [{:03d}/{:03d}] loss={:.4f}".format(
|
||||
idx, len(env), future_loss.item()
|
||||
)
|
||||
)
|
||||
meta_model.adapt(
|
||||
future_time,
|
||||
future_x,
|
||||
future_y,
|
||||
env.timestamp_interval,
|
||||
args.refine_lr,
|
||||
args.refine_epochs,
|
||||
)
|
||||
import pdb
|
||||
|
||||
pdb.set_trace()
|
||||
for iseq in range(args.seq_length):
|
||||
time_seqs.append(future_time - iseq * eval_env.timestamp_interval)
|
||||
print("-")
|
||||
|
||||
|
||||
@ -156,6 +152,7 @@ def pretrain_v2(base_model, meta_model, criterion, xenv, args, logger):
|
||||
meta_model.set_best_name("pretrain-{:}.pth".format(args.rand_seed))
|
||||
last_success_epoch, early_stop_thresh = 0, args.pretrain_early_stop_thresh
|
||||
per_epoch_time, start_time = AverageMeter(), time.time()
|
||||
device = args.device
|
||||
for iepoch in range(args.epochs):
|
||||
left_time = "Time Left: {:}".format(
|
||||
convert_secs2time(per_epoch_time.avg * (args.epochs - iepoch), True)
|
||||
@ -163,32 +160,38 @@ def pretrain_v2(base_model, meta_model, criterion, xenv, args, logger):
|
||||
total_meta_v1_losses, total_meta_v2_losses, total_match_losses = [], [], []
|
||||
optimizer.zero_grad()
|
||||
for ibatch in range(args.meta_batch):
|
||||
rand_index = random.randint(0, meta_model.meta_length - 1)
|
||||
timestamp = meta_model.meta_timestamps[rand_index]
|
||||
meta_embed = meta_model.super_meta_embed[rand_index]
|
||||
|
||||
timestamps, [container], time_embeds = meta_model(
|
||||
torch.unsqueeze(timestamp, dim=0), None, True
|
||||
)
|
||||
_, (inputs, targets) = xenv(timestamp.item())
|
||||
inputs, targets = inputs.to(device), targets.to(device)
|
||||
# generate models one step ahead
|
||||
predictions = base_model.forward_with_container(inputs, container)
|
||||
total_meta_v1_losses.append(criterion(predictions, targets))
|
||||
# the matching loss
|
||||
match_loss = criterion(torch.squeeze(time_embeds, dim=0), meta_embed)
|
||||
total_match_losses.append(match_loss)
|
||||
# generate models via memory
|
||||
rand_index = random.randint(0, meta_model.meta_length - xenv.seq_length - 1)
|
||||
_, [seq_containers], _ = meta_model(
|
||||
None,
|
||||
torch.unsqueeze(
|
||||
meta_model.super_meta_embed[
|
||||
rand_index : rand_index + xenv.seq_length
|
||||
],
|
||||
dim=0,
|
||||
),
|
||||
False,
|
||||
)
|
||||
timestamps = meta_model.meta_timestamps[
|
||||
rand_index : rand_index + xenv.seq_length
|
||||
]
|
||||
meta_embeds = meta_model.super_meta_embed[
|
||||
rand_index : rand_index + xenv.seq_length
|
||||
]
|
||||
|
||||
_, (seq_inputs, seq_targets) = xenv.seq_call(timestamps)
|
||||
seq_inputs, seq_targets = seq_inputs.to(args.device), seq_targets.to(
|
||||
args.device
|
||||
)
|
||||
# generate models one step ahead
|
||||
[seq_containers], time_embeds = meta_model(
|
||||
torch.unsqueeze(timestamps, dim=0), None
|
||||
)
|
||||
for container, inputs, targets in zip(
|
||||
seq_containers, seq_inputs, seq_targets
|
||||
):
|
||||
predictions = base_model.forward_with_container(inputs, container)
|
||||
total_meta_v1_losses.append(criterion(predictions, targets))
|
||||
# the matching loss
|
||||
match_loss = criterion(torch.squeeze(time_embeds, dim=0), meta_embeds)
|
||||
total_match_losses.append(match_loss)
|
||||
# generate models via memory
|
||||
[seq_containers], _ = meta_model(None, torch.unsqueeze(meta_embeds, dim=0))
|
||||
seq_inputs, seq_targets = seq_inputs.to(device), seq_targets.to(device)
|
||||
for container, inputs, targets in zip(
|
||||
seq_containers, seq_inputs, seq_targets
|
||||
):
|
||||
@ -250,7 +253,14 @@ def main(args):
|
||||
|
||||
# pre-train the hypernetwork
|
||||
timestamps = train_env.get_timestamp(None)
|
||||
meta_model = LFNA_Meta(shape_container, args.layer_dim, args.time_dim, timestamps)
|
||||
meta_model = LFNA_Meta(
|
||||
shape_container,
|
||||
args.layer_dim,
|
||||
args.time_dim,
|
||||
timestamps,
|
||||
seq_length=args.seq_length,
|
||||
interval=train_env.timestamp_interval,
|
||||
)
|
||||
meta_model = meta_model.to(args.device)
|
||||
|
||||
logger.log("The base-model has {:} weights.".format(base_model.numel()))
|
||||
|
||||
@ -22,7 +22,9 @@ class LFNA_Meta(super_core.SuperModule):
|
||||
meta_timestamps,
|
||||
mha_depth: int = 2,
|
||||
dropout: float = 0.1,
|
||||
thresh: float = 0.05,
|
||||
seq_length: int = 10,
|
||||
interval: float = None,
|
||||
thresh: float = None,
|
||||
):
|
||||
super(LFNA_Meta, self).__init__()
|
||||
self._shape_container = shape_container
|
||||
@ -31,7 +33,10 @@ class LFNA_Meta(super_core.SuperModule):
|
||||
for ilayer in range(self._num_layers):
|
||||
self._numel_per_layer.append(shape_container[ilayer].numel())
|
||||
self._raw_meta_timestamps = meta_timestamps
|
||||
self._thresh = thresh
|
||||
assert interval is not None
|
||||
self._interval = interval
|
||||
self._seq_length = seq_length
|
||||
self._thresh = interval * 30 if thresh is None else thresh
|
||||
|
||||
self.register_parameter(
|
||||
"_super_layer_embed",
|
||||
@ -42,6 +47,10 @@ class LFNA_Meta(super_core.SuperModule):
|
||||
torch.nn.Parameter(torch.Tensor(len(meta_timestamps), time_embedding)),
|
||||
)
|
||||
self.register_buffer("_meta_timestamps", torch.Tensor(meta_timestamps))
|
||||
# register a time difference buffer
|
||||
time_interval = [-i * self._interval for i in range(self._seq_length)]
|
||||
time_interval.reverse()
|
||||
self.register_buffer("_time_interval", torch.Tensor(time_interval))
|
||||
self._time_embed_dim = time_embedding
|
||||
self._append_meta_embed = dict(fixed=None, learnt=None)
|
||||
self._append_meta_timestamps = dict(fixed=None, learnt=None)
|
||||
@ -51,12 +60,12 @@ class LFNA_Meta(super_core.SuperModule):
|
||||
)
|
||||
|
||||
# build transformer
|
||||
self._trans_att = super_core.SuperQKVAttentionV2(
|
||||
qk_att_dim=time_embedding,
|
||||
in_v_dim=time_embedding,
|
||||
hidden_dim=time_embedding,
|
||||
self._trans_att = super_core.SuperQKVAttention(
|
||||
time_embedding,
|
||||
time_embedding,
|
||||
time_embedding,
|
||||
time_embedding,
|
||||
num_heads=4,
|
||||
proj_dim=time_embedding,
|
||||
qkv_bias=True,
|
||||
attn_drop=None,
|
||||
proj_drop=dropout,
|
||||
@ -166,12 +175,9 @@ class LFNA_Meta(super_core.SuperModule):
|
||||
# timestamps is a batch of sequence of timestamps
|
||||
batch, seq = timestamps.shape
|
||||
meta_timestamps, meta_embeds = self.meta_timestamps, self.super_meta_embed
|
||||
# timestamp_q_embed = self._tscalar_embed(timestamps)
|
||||
# timestamp_k_embed = self._tscalar_embed(meta_timestamps.view(1, -1))
|
||||
timestamp_q_embed = self._tscalar_embed(timestamps)
|
||||
timestamp_k_embed = self._tscalar_embed(meta_timestamps.view(1, -1))
|
||||
timestamp_v_embed = meta_embeds.unsqueeze(dim=0)
|
||||
timestamp_qk_att_embed = self._tscalar_embed(
|
||||
torch.unsqueeze(timestamps, dim=-1) - meta_timestamps
|
||||
)
|
||||
# create the mask
|
||||
mask = (
|
||||
torch.unsqueeze(timestamps, dim=-1) <= meta_timestamps.view(1, 1, -1)
|
||||
@ -182,7 +188,7 @@ class LFNA_Meta(super_core.SuperModule):
|
||||
> self._thresh
|
||||
)
|
||||
timestamp_embeds = self._trans_att(
|
||||
timestamp_qk_att_embed, timestamp_v_embed, mask
|
||||
timestamp_q_embed, timestamp_k_embed, timestamp_v_embed, mask
|
||||
)
|
||||
relative_timestamps = timestamps - timestamps[:, :1]
|
||||
relative_pos_embeds = self._tscalar_embed(relative_timestamps)
|
||||
@ -192,36 +198,69 @@ class LFNA_Meta(super_core.SuperModule):
|
||||
corrected_embeds = self._meta_corrector(init_timestamp_embeds)
|
||||
return corrected_embeds
|
||||
|
||||
def forward_raw(self, timestamps, time_embed):
|
||||
if time_embed is None:
|
||||
batch, seq = timestamps.shape
|
||||
time_embed = self._obtain_time_embed(timestamps)
|
||||
def forward_raw(self, timestamps, time_embeds, get_seq_last):
|
||||
if time_embeds is None:
|
||||
time_seq = timestamps.view(-1, 1) + self._time_interval.view(1, -1)
|
||||
B, S = time_seq.shape
|
||||
time_embeds = self._obtain_time_embed(time_seq)
|
||||
else:
|
||||
batch, seq, _ = time_embed.shape
|
||||
time_seq = None
|
||||
B, S, _ = time_embeds.shape
|
||||
# create joint embed
|
||||
num_layer, _ = self._super_layer_embed.shape
|
||||
meta_embed = time_embed.view(batch, seq, 1, -1).expand(-1, -1, num_layer, -1)
|
||||
layer_embed = self._super_layer_embed.view(1, 1, num_layer, -1).expand(
|
||||
batch, seq, -1, -1
|
||||
)
|
||||
joint_embed = torch.cat(
|
||||
(meta_embed, layer_embed), dim=-1
|
||||
) # batch, seq, num-layers, input-dim
|
||||
batch_weights = self._generator(
|
||||
joint_embed
|
||||
) # batch, seq, num-layers, num-weights
|
||||
if get_seq_last:
|
||||
time_embeds = time_embeds[:, -1, :]
|
||||
# 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,
|
||||
)
|
||||
else:
|
||||
# The shape of `joint_embed` is batch * seq * num-layers * input-dim
|
||||
joint_embeds = torch.cat(
|
||||
(
|
||||
time_embeds.view(B, S, 1, -1).expand(-1, -1, num_layer, -1),
|
||||
self._super_layer_embed.view(1, 1, num_layer, -1).expand(
|
||||
B, S, -1, -1
|
||||
),
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
batch_weights = self._generator(joint_embeds)
|
||||
batch_containers = []
|
||||
for seq_weights in torch.split(batch_weights, 1):
|
||||
seq_containers = []
|
||||
for weights in torch.split(seq_weights.squeeze(0), 1):
|
||||
weights = torch.split(weights.squeeze(0), 1)
|
||||
seq_containers.append(self._shape_container.translate(weights))
|
||||
batch_containers.append(seq_containers)
|
||||
return batch_containers, time_embed
|
||||
for weights in torch.split(batch_weights, 1):
|
||||
if get_seq_last:
|
||||
batch_containers.append(
|
||||
self._shape_container.translate(torch.split(weights.squeeze(0), 1))
|
||||
)
|
||||
else:
|
||||
seq_containers = []
|
||||
for ws in torch.split(weights.squeeze(0), 1):
|
||||
seq_containers.append(
|
||||
self._shape_container.translate(torch.split(ws.squeeze(0), 1))
|
||||
)
|
||||
batch_containers.append(seq_containers)
|
||||
return time_seq, batch_containers, time_embeds
|
||||
|
||||
def forward_candidate(self, input):
|
||||
raise NotImplementedError
|
||||
|
||||
def adapt(self, timestamp, x, y, threshold, lr, epochs):
|
||||
if distance + threshold * 1e-2 <= threshold:
|
||||
return False
|
||||
with torch.set_grad_enabled(True):
|
||||
new_param = self.create_meta_embed()
|
||||
optimizer = torch.optim.Adam(
|
||||
[new_param], lr=args.refine_lr, weight_decay=1e-5, amsgrad=True
|
||||
)
|
||||
import pdb
|
||||
|
||||
pdb.set_trace()
|
||||
print("-")
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return "(_super_layer_embed): {:}, (_super_meta_embed): {:}, (_meta_timestamps): {:}".format(
|
||||
list(self._super_layer_embed.shape),
|
||||
|
||||
Loading…
Reference in New Issue
Block a user