wassname/attentive-neural-processes
1
0
Fork 0
mirror of https://github.com/wassname/attentive-neural-processes.git synced 2026-08-02 12:30:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor

Browse Source
This commit is contained in:
wassname
2020-04-19 11:49:59 +08:00
parent 5af968080f
commit db554df064
10 changed files with 445 additions and 1087 deletions
Download Patch File Download Diff File Expand all files Collapse all files
neural_processes/lightning.py
+19 -33
View File
@@ -35,40 +35,27 @@ class PL_Seq2Seq(pl.LightningModule):
assert all(torch.isfinite(d).all() for d in batch)
context_x, context_y, target_x, target_y = batch
y_dist, losses, extra = self.forward(context_x, context_y, target_x, target_y)
loss = losses['loss_p'] # + loss_mse
tensorboard_logs = {
"train/loss": loss,
'train/loss_mse': losses['loss_mse'],
"train/loss_p": losses['loss_p'],
"train/sigma": torch.exp(extra['log_sigma']).mean()}
return {"loss": loss, "log": tensorboard_logs}
tensorboard_logs = {"train_" + k: v for k, v in losses.items()}
assert torch.isfinite(tensorboard_logs["train_loss"])
return {"loss": tensorboard_logs['train_loss'], "log": tensorboard_logs}
def validation_step(self, batch, batch_idx):
context_x, context_y, target_x, target_y = batch
assert all(torch.isfinite(d).all() for d in batch)
y_dist, losses, extra = self.forward(context_x, context_y, target_x, target_y)
loss = losses['loss_p'] # + loss_mse
tensorboard_logs = {
"val_loss": loss,
'val/loss_mse': losses['loss_mse'],
"val/loss_p": losses['loss_p'],
"val/sigma": torch.exp(extra['log_sigma']).mean()}
return {"val_loss": loss, "log": tensorboard_logs}
tensorboard_logs = {"val_" + k: v for k, v in losses.items()}
assert torch.isfinite(tensorboard_logs["val_loss"])
return {"val_loss": tensorboard_logs["val_loss"], "log": tensorboard_logs}
def validation_end(self, outputs):
if int(self.hparams["vis_i"]) > 0:
self.show_image()
avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
keys = outputs[0]["log"].keys()
tensorboard_logs = {
k: torch.stack([x["log"][k] for x in outputs if k in x["log"]]).mean()
for k in keys
}
tensorboard_logs = self.agg_logs(outputs)
tensorboard_logs_str = {k: f"{v}" for k, v in tensorboard_logs.items()}
print(f"step {self.trainer.global_step}, {tensorboard_logs_str}")
assert torch.isfinite(avg_loss)
return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
return {"avg_val_loss": tensorboard_logs["val_loss"], "log": tensorboard_logs}
def agg_logs(self, outputs):
if isinstance(outputs, dict):
@@ -99,25 +86,24 @@ class PL_Seq2Seq(pl.LightningModule):
def test_step(self, batch, batch_idx):
pred, losses, extra = self.forward(*batch)
# For test use a diff loss, MSE over next 24
# loss = losses["loss"]
loss = F.mse_loss(pred, batch[-1], reduction='none')[:, :24].mean()
context_x, context_y, target_x, target_y = batch
y_dist = extra['y_dist']
# For test use a diff loss, log_p over next <24h, so it's a standard amount of steps
loss = -y_dist.log_prob(target_y)[:, :24].mean()
tensorboard_logs = {"test_" + k: v for k, v in losses.items()}
assert torch.isfinite(loss)
return {"test_loss": loss, "log": tensorboard_logs}
def test_end(self, outputs):
avg_loss = torch.stack([x["test_loss"] for x in outputs]).mean()
keys = outputs[0]["log"].keys()
tensorboard_logs = {
k: torch.stack([x["log"][k] for x in outputs if k in x["log"]]).mean()
for k in keys
}
tensorboard_logs_str = {k: f"{v}" for k, v in tensorboard_logs.items()}
tensorboard_logs = self.agg_logs(outputs)
logger.info(
f"step {self.trainer.global_step}, {tensorboard_logs_str}"
)
return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
return {"avg_test_loss": tensorboard_logs["test_loss"], "log": tensorboard_logs}
def configure_optimizers(self):
optim = torch.optim.Adam(self.parameters(), lr=self.hparams["learning_rate"])
neural_processes/models/lstm_seqseq.py
+7 -6
View File
@@ -109,7 +109,7 @@ class Seq2SeqNet(nn.Module):
y_dist = torch.distributions.Normal(mean, sigma)
# Loss
loss_mse = loss_p = None
loss_mse = loss_p = loss_p_weighted = None
if target_y is not None:
loss_mse = F.mse_loss(mean, target_y, reduction="none")
if self._use_lvar:
@@ -120,15 +120,16 @@ class Seq2SeqNet(nn.Module):
if self.hparams["context_in_target"]:
loss_p[: context_x.size(1)] /= 100
loss_mse[: context_x.size(1)] /= 100
# # Don't catch loss on context window
# mean = mean[:, self.hparams.num_context:]
# log_sigma = log_sigma[:, self.hparams.num_context:]
# Weight loss nearer to prediction time?
weight = (torch.arange(loss_p.shape[1]) + 1).float().to(device)[None, :]
loss_p_weighted = loss_p / torch.sqrt(weight)
y_pred = y_dist.rsample if self.training else y_dist.loc
return (
y_pred,
dict(loss_p=loss_p.mean(), loss_mse=loss_mse.mean()),
dict(log_sigma=log_sigma, dist=y_dist),
dict(loss=loss_p.mean(), loss_p=loss_p.mean(), loss_mse=loss_mse.mean(), loss_p_weighted=loss_p_weighted.mean()),
dict(log_sigma=log_sigma, y_dist=y_dist),
)
neural_processes/models/lstm_std.py
+4 -7
View File
@@ -76,7 +76,7 @@ class LSTMNet(nn.Module):
y_dist = torch.distributions.Normal(mean, sigma)
# Loss
loss_mse = loss_p = None
loss_mse = loss_p_weighted = loss_p = None
if target_y is not None:
loss_mse = F.mse_loss(mean, target_y, reduction="none")
if self._use_lvar:
@@ -86,19 +86,16 @@ class LSTMNet(nn.Module):
if self.hparams["context_in_target"]:
loss_p[: context_x.size(1)] /= 100
loss_mse[: context_x.size(1)] /= 100
# # Don't catch loss on context window
# mean = mean[:, self.hparams.num_context:]
# log_sigma = log_sigma[:, self.hparams.num_context:]
# Weight loss nearer to prediction time?
weight = (torch.arange(loss_p.shape[1]) + 1).float().to(device)[None, :]
loss_p = loss_p / torch.sqrt(weight) # We want to weight nearer stuff more
loss_p_weighted = loss_p / torch.sqrt(weight) # We want to weight nearer stuff more
y_pred = y_dist.rsample if self.training else y_dist.loc
return (
y_pred,
dict(loss_p=loss_p.mean(), loss_mse=loss_mse.mean()),
dict(log_sigma=log_sigma, dist=y_dist),
dict(loss=loss_p.mean(), loss_p_weighted=loss_p_weighted.mean(), loss_p=loss_p.mean(), loss_mse=loss_mse.mean()),
dict(log_sigma=log_sigma, y_dist=y_dist),
)
# loss = None
# if target_y is not None:
neural_processes/models/neural_process/lightning.py
+175 -35
View File
@@ -9,32 +9,50 @@ from neural_processes.utils import ObjectDict
class PL_NeuralProcess(PL_Seq2Seq):
"""Base class with everything off."""
def __init__(self, hparams,
MODEL_CLS=NeuralProcess.FROM_HPARAMS, **kwargs):
super().__init__(hparams,
MODEL_CLS=MODEL_CLS, **kwargs)
DEFAULT_ARGS = {
'attention_dropout': 0,
'attention_layers': 2,
'batchnorm': False,
'det_enc_cross_attn_type': 'multihead',
'det_enc_self_attn_type': 'uniform',
'dropout': 0,
'hidden_dim_power': 7,
'latent_dim_power': 7,
'learning_rate': 0.006,
'attention_dropout': 0,
'batchnorm': False,
'attention_layers': 2,
'det_enc_cross_attn_type': 'uniform',
'det_enc_self_attn_type': 'uniform',
'latent_enc_self_attn_type': 'uniform',
'learning_rate': 0.002,
'n_decoder_layers': 4,
'n_det_encoder_layers': 4,
'n_latent_encoder_layers': 2,
'num_heads_power': 3,
'use_deterministic_path': True,
'use_lvar': True,
'use_self_attn': True,
'hidden_dim_power': 5,
'latent_dim_power': 4,
'n_decoder_layers': 4,
'n_latent_encoder_layers': 2,
'use_deterministic_path': False,
'n_det_encoder_layers': 4,
'use_lvar': False,
'use_self_attn': False,
'use_rnn': False,
}
USR_ATTRS_DEFAULT = {
'batch_size': 16,
'grad_clip': 40,
'max_nb_epochs': 200,
'num_workers': 4,
'num_context': 24* 4,
'vis_i': '670',
'num_extra_target': 24*4,
'x_dim': 18,
'context_in_target': False,
'y_dim': 1,
'patience': 3,
'min_std': 0.005,
}
@staticmethod
def add_suggest(trial, user_attrs={}):
trial.suggest_loguniform("learning_rate", 1e-6, 1e-2)
@@ -42,15 +60,11 @@ class PL_NeuralProcess(PL_Seq2Seq):
trial.suggest_discrete_uniform("num_heads_power", 2, 4, 1)
trial.suggest_discrete_uniform(
"hidden_dim_power", 3, 11, 1
"hidden_dim_power", 4, 11, 1
)
trial.suggest_discrete_uniform(
"latent_dim_power", 3, 11, 1
"latent_dim_power", 4, 11, 1
)
trial.suggest_int(
"n_latent_encoder_layers", 1, 11
)
trial.suggest_int("n_latent_encoder_layers", 1, 12)
trial.suggest_int("n_det_encoder_layers", 1, 12)
trial.suggest_int("n_decoder_layers", 1, 12)
@@ -70,21 +84,147 @@ class PL_NeuralProcess(PL_Seq2Seq):
trial.suggest_categorical("use_deterministic_path", [False, True])
trial.suggest_categorical("use_rnn", [True, False])
user_attrs_default = {
'batch_size': 16,
'grad_clip': 40,
'max_nb_epochs': 200,
'num_workers': 4,
'num_context': 24* 4,
'vis_i': '670',
'num_extra_target': 24*4,
'x_dim': 18,
'context_in_target': False,
'y_dim': 1,
'patience': 3,
'min_std': 0.005,
}
[trial.set_user_attr(k, v) for k, v in user_attrs_default.items()]
[trial.set_user_attr(k, v) for k, v in PL_NeuralProcess.USR_ATTRS_DEFAULT.items()]
[trial.set_user_attr(k, v) for k, v in user_attrs.items()]
return trial
class PL_NP(PL_NeuralProcess):
"""Vanilla NP with no attention or RNN."""
def __init__(self, hparams,
MODEL_CLS=NeuralProcess.FROM_HPARAMS, **kwargs):
super().__init__(hparams,
MODEL_CLS=MODEL_CLS, **kwargs)
DEFAULT_ARGS = {
**PL_NeuralProcess.DEFAULT_ARGS,
'det_enc_cross_attn_type': 'uniform',
'det_enc_self_attn_type': 'uniform',
'latent_enc_self_attn_type': 'uniform',
'use_deterministic_path': False,
}
@staticmethod
def add_suggest(trial, user_attrs={}):
trial.suggest_loguniform("learning_rate", 1e-6, 1e-2)
trial.suggest_discrete_uniform(
"hidden_dim_power", 3, 11, 1
)
trial.suggest_discrete_uniform(
"latent_dim_power", 3, 11, 1
)
trial.suggest_int("n_latent_encoder_layers", 1, 12)
trial.suggest_int("n_decoder_layers", 1, 12)
trial.suggest_uniform("dropout", 0, 0.9)
trial.suggest_categorical("batchnorm", [False, True])
[trial.set_user_attr(k, v) for k, v in PL_NeuralProcess.USR_ATTRS_DEFAULT.items()]
[trial.set_user_attr(k, v) for k, v in user_attrs.items()]
return trial
class PL_ANP(PL_NeuralProcess):
def __init__(self, hparams,
MODEL_CLS=NeuralProcess.FROM_HPARAMS, **kwargs):
super().__init__(hparams,
MODEL_CLS=MODEL_CLS, **kwargs)
DEFAULT_ARGS = {
**PL_NeuralProcess.DEFAULT_ARGS,
'det_enc_cross_attn_type': 'multihead',
'det_enc_self_attn_type': 'multihead',
'latent_enc_self_attn_type': 'multihead',
'use_self_attn': True,
'use_deterministic_path': True,
}
@staticmethod
def add_suggest(trial, user_attrs={}):
trial.suggest_loguniform("learning_rate", 1e-6, 1e-2)
trial.suggest_int("attention_layers", 1, 4)
trial.suggest_discrete_uniform("num_heads_power", 2, 4, 1)
trial.suggest_discrete_uniform(
"hidden_dim_power", 4, 11, 1
)
trial.suggest_discrete_uniform(
"latent_dim_power", 4, 11, 1
)
trial.suggest_int("n_latent_encoder_layers", 1, 12)
trial.suggest_int("n_det_encoder_layers", 1, 12)
trial.suggest_int("n_decoder_layers", 1, 12)
trial.suggest_uniform("dropout", 0, 0.9)
trial.suggest_uniform("attention_dropout", 0, 0.9)
trial.suggest_categorical(
"latent_enc_self_attn_type", ['uniform', 'multihead']
)
trial.suggest_categorical("det_enc_self_attn_type", ['uniform', 'multihead'])
trial.suggest_categorical("det_enc_cross_attn_type", ['uniform', 'multihead'])
trial.suggest_categorical("batchnorm", [False, True])
trial.suggest_categorical("use_deterministic_path", [False, True])
[trial.set_user_attr(k, v) for k, v in PL_NeuralProcess.USR_ATTRS_DEFAULT.items()]
[trial.set_user_attr(k, v) for k, v in user_attrs.items()]
return trial
class PL_ANPRNN(PL_NeuralProcess):
"""
Recurrent Attentive Neural Process for Sequential Data.
https://arxiv.org/abs/1910.09323
"""
def __init__(self, hparams,
MODEL_CLS=NeuralProcess.FROM_HPARAMS, **kwargs):
super().__init__(hparams,
MODEL_CLS=MODEL_CLS, **kwargs)
DEFAULT_ARGS = {
**PL_NeuralProcess.DEFAULT_ARGS,
'det_enc_cross_attn_type': 'multihead',
'det_enc_self_attn_type': 'multihead',
'latent_enc_self_attn_type': 'multihead',
'use_self_attn': True,
'use_rnn': True,
}
@staticmethod
def add_suggest(trial, user_attrs={}):
trial.suggest_loguniform("learning_rate", 1e-6, 1e-2)
trial.suggest_int("attention_layers", 1, 4)
trial.suggest_discrete_uniform("num_heads_power", 2, 4, 1)
trial.suggest_discrete_uniform(
"hidden_dim_power", 4, 11, 1
)
trial.suggest_discrete_uniform(
"latent_dim_power", 4, 11, 1
)
trial.suggest_int("n_latent_encoder_layers", 1, 12)
trial.suggest_int("n_det_encoder_layers", 1, 12)
trial.suggest_int("n_decoder_layers", 1, 12)
trial.suggest_uniform("dropout", 0, 0.9)
trial.suggest_uniform("attention_dropout", 0, 0.9)
trial.suggest_categorical("batchnorm", [False, True])
trial.suggest_categorical("use_deterministic_path", [False, True])
[trial.set_user_attr(k, v) for k, v in PL_NeuralProcess.USR_ATTRS_DEFAULT.items()]
[trial.set_user_attr(k, v) for k, v in user_attrs.items()]
return trial
neural_processes/models/neural_process/model.py
+13 -3
View File
@@ -304,7 +304,8 @@ class NeuralProcess(nn.Module):
self._use_lvar = use_lvar
def forward(self, context_x, context_y, target_x, target_y=None):
device = next(self.parameters()).device
# https://stackoverflow.com/a/46772183/221742
target_x = self.norm_x(target_x)
context_x = self.norm_x(context_x)
@@ -353,18 +354,27 @@ class NeuralProcess(nn.Module):
log_p[:, :context_x.size(1)] /= 100 # There's the temptation for it to fit only on context, where it knows the answer, and learn very low uncertainty.
loss_kl = torch.distributions.kl_divergence(
dist_post, dist_prior).mean(-1) # [B, R].mean(-1)
loss_kl = loss_kl[:, None].expand(log_p.shape)
mse_loss = F.mse_loss(dist.loc, target_y, reduction='none')[:,:context_x.size(1)].mean()
loss_p = -log_p.mean()
loss_p = -log_p
# Weight loss nearer to prediction time?
weight = (torch.arange(loss_p.shape[1]) + 1).float().to(device)[None, :]
loss_p_weighted = loss_p / torch.sqrt(weight) # We want to weight nearer stuff more
loss_p_weighted = loss_p_weighted.mean()
loss = (loss_kl - log_p).mean()
loss_kl = loss_kl.mean()
log_p = log_p.mean()
loss_p = loss_p.mean()
else:
loss_p = None
mse_loss = None
loss_kl = None
loss = None
loss_p_weighted = None
y_pred = dist.rsample() if self.training else dist.loc
return y_pred, dict(loss=loss, loss_p=loss_p, loss_kl=loss_kl, loss_mse=mse_loss), dict(log_sigma=log_sigma, dist=dist)
return y_pred, dict(loss=loss, loss_p=loss_p, loss_kl=loss_kl, loss_mse=mse_loss, loss_p_weighted=loss_p_weighted), dict(log_sigma=log_sigma, dist=dist)
neural_processes/models/transformer.py
+5 -74
View File
@@ -68,22 +68,17 @@ class NetTransformer(nn.Module):
x[~x_mask] = 0
x = x.detach()
x_key_padding_mask = ~x_mask.any(-1)
# print('x_key_padding_mask', x_mask.float().mean())
# print(x.shape, 'x1')
x = self.enc_emb(x).permute(1, 0, 2)
# print(x.shape, 'x2')
# Size([C, B, emb_dim])
outputs = self.encoder(x, src_key_padding_mask=x_key_padding_mask).permute(
1, 0, 2
)
# print(outputs.shape, 'outputs')
# Seems to help a little, especially with extrapolating out of bounds
steps = context_y.shape[1]
mean = self.mean(outputs)[:, steps:, :]
log_sigma = self.std(outputs)[:, steps:, :]
# mean_target = mean[:, -steps:, :]
# mean_context = mean[:, :-steps, :]
if self._use_lvar:
log_sigma = torch.clamp(
@@ -105,81 +100,17 @@ class NetTransformer(nn.Module):
if self.hparams["context_in_target"]:
loss_p[: context_x.size(1)] /= 100
loss_mse[: context_x.size(1)] /= 100
# # Don't catch loss on context window
# mean = mean[:, self.hparams.num_context:]
# log_sigma = log_sigma[:, self.hparams.num_context:]
# Weight loss nearer to prediction time?
weight = (torch.arange(loss_p.shape[1]) + 1).float().to(device)[None, :]
loss_p = loss_p / torch.sqrt(weight) # We want to weight nearer stuff more
loss_p_weighted = loss_p / torch.sqrt(weight) # We want to weight nearer stuff more
y_pred = y_dist.rsample if self.training else y_dist.loc
return (
y_pred,
dict(loss_p=loss_p.mean(), loss_mse=loss_mse.mean()),
dict(log_sigma=log_sigma, dist=y_dist),
dict(loss=loss_p.mean(), loss_p=loss_p.mean(), loss_mse=loss_mse.mean(), loss_p_weighted=loss_p_weighted.mean()),
dict(log_sigma=log_sigma, y_dist=y_dist),
)
# mean_target = mean[:, -steps:, :]
# mean_context = mean[:, :-steps, :]
# loss = None
# if target_y is not None:
# y = torch.cat([context_y, target_y], 1)
# y_mask = torch.isfinite(y) & (y != self.hparams.nan_value)
# y[~y_mask] = 0
# y = y.detach()
# loss_scale = 100
# # loss = F.mse_loss(mean * loss_scale, y * loss_scale, reduction='none') / loss_scale
# loss_target = (
# F.mse_loss(
# mean_target * loss_scale,
# y[:, -steps:, :] * loss_scale,
# reduction="none",
# )
# / loss_scale
# )
# loss_context = (
# F.mse_loss(
# mean_context * loss_scale,
# y[:, :-steps, :] * loss_scale,
# reduction="none",
# )
# / loss_scale
# )
# y_mask_target = y_mask[:, -steps:, :].detach()
# y_mask_context = y_mask[:, :-steps, :].detach()
# # loss_target = loss[:, -steps:, :]
# # loss_context = loss[:, :-steps, :]
# # print(0, loss_context.sum(), loss_target.sum())
# weight = (
# (torch.arange(loss_target.shape[1]) + 0.5)
# .float()
# .to(device)[None, :, None]
# )
# # weight /= weight.sum()
# # print(1.0, loss_context.sum(), loss_target.sum())
# loss_target = loss_target / torch.sqrt(
# weight
# ) # We want to weight nearer stuff more
# # print(1.5, loss_context.sum(), y_mask_context.sum(), loss_target.sum(), y_mask_target.sum(), (loss_context * y_mask_context).sum())
# loss_context = (loss_context * y_mask_context.float()).sum() / (
# y_mask_context.sum() + 1.0
# )
# loss_target = (loss_target * y_mask_target.float()).sum() / (
# y_mask_target.sum() + 1.0
# ) # Mean over unmasked ones
# # print(2, loss_context.sum(), loss_target.sum())
# # Perhaps predicting the past, as a secondary loss will help
# loss = loss_context / 100.0 + loss_target
# assert torch.isfinite(loss)
# return mean_target, dict(loss=loss), dict()
class PL_Transformer(PL_Seq2Seq):
neural_processes/models/transformer_seq2seq.py
+14 -15
View File
@@ -110,7 +110,7 @@ class TransformerSeq2SeqNet(nn.Module):
y_dist = torch.distributions.Normal(mean, sigma)
# Loss
loss_mse = loss_p = None
loss_mse = loss_p = loss_p_weighted = None
if target_y is not None:
loss_mse = F.mse_loss(mean, target_y, reduction="none")
if self._use_lvar:
@@ -120,19 +120,16 @@ class TransformerSeq2SeqNet(nn.Module):
if self.hparams["context_in_target"]:
loss_p[: context_x.size(1)] /= 100
loss_mse[: context_x.size(1)] /= 100
# # Don't catch loss on context window
# mean = mean[:, self.hparams.num_context:]
# log_sigma = log_sigma[:, self.hparams.num_context:]
# Weight loss nearer to prediction time?
weight = (torch.arange(loss_p.shape[1]) + 1).float().to(device)[None, :]
loss_p = loss_p / torch.sqrt(weight) # We want to weight nearer stuff more
loss_p_weighted = loss_p / torch.sqrt(weight) # We want to weight nearer stuff more
y_pred = y_dist.rsample if self.training else y_dist.loc
return (
y_pred,
dict(loss_p=loss_p.mean(), loss_mse=loss_mse.mean()),
dict(log_sigma=log_sigma, dist=y_dist),
dict(loss=loss_p.mean(), loss_p=loss_p.mean(), loss_mse=loss_mse.mean(), loss_p_weighted=loss_p_weighted.mean()),
dict(log_sigma=log_sigma, y_dist=y_dist),
)
@@ -141,13 +138,13 @@ class TransformerSeq2Seq_PL(PL_Seq2Seq):
super().__init__(hparams, MODEL_CLS=MODEL_CLS, **kwargs)
DEFAULT_ARGS = {
"agg": "mean",
"attention_dropout": 0.12,
"agg": "max",
"attention_dropout": 0.2,
"hidden_out_size_power": 4,
"hidden_size_power": 7,
"learning_rate": 0.0023,
"nhead_power": 2,
"nlayers": 4,
"hidden_size_power": 5,
"learning_rate": 0.006,
"nhead_power": 3,
"nlayers": 2,
}
@staticmethod
@@ -169,8 +166,10 @@ class TransformerSeq2Seq_PL(PL_Seq2Seq):
"""
trial.suggest_loguniform("learning_rate", 1e-6, 1e-2)
trial.suggest_uniform("attention_dropout", 0, 0.75)
trial.suggest_discrete_uniform("hidden_size_power", 2, 10, 1)
trial.suggest_discrete_uniform("hidden_out_size_power", 2, 9, 1)
# we must have nhead<==hidden_size
# so nhead_power.max()<==hidden_size_power.min()
trial.suggest_discrete_uniform("hidden_size_power", 4, 10, 1)
trial.suggest_discrete_uniform("hidden_out_size_power", 4, 9, 1)
trial.suggest_discrete_uniform("nhead_power", 1, 4, 1)
trial.suggest_int("nlayers", 1, 12)
neural_processes/train.py
+18 -6
View File
@@ -56,21 +56,32 @@ def main(
return model, trainer
def objective(trial, PL_MODEL_CLS, name):
def objective(trial, PL_MODEL_CLS, name, user_attrs):
"""For optuna hparam opt."""
# see https://github.com/optuna/optuna/blob/cf6f02d/examples/pytorch_lightning_simple.py
trial = PL_MODEL_CLS.add_suggest(trial)
# trial._user_attrs.update(user_attrs)
[trial.set_user_attr(k, v) for k, v in user_attrs.items()]
print("trial", trial.number, "params", trial.params, trial._user_attrs)
print(dict(number=trial.number, params=trial.params, user_attrs=trial.user_attrs))
model, trainer = main(trial, PL_MODEL_CLS=PL_MODEL_CLS, name=name)
# Load checkpoint
checkpoints = sorted(Path(trainer.checkpoint_callback.dirpath).glob("*.ckpt"))
if len(checkpoints):
checkpoint = checkpoints[-1]
device = next(model.parameters()).device
print(f"Loading checkpoint {checkpoint}")
model = model.load_from_checkpoint(checkpoint).to(device)
trainer.test(model)
# also report to tensorboard & print
print("logger.metrics", model.logger.metrics[-1:])
model.logger.experiment.add_hparams(trial.params, logger.metrics[-1])
model.logger.experiment.add_hparams(trial.params, model.logger.metrics[-1])
model.logger.save()
return model.logger.metrics[-1]["val_loss"]
return model.logger.metrics[-1]["avg_test_loss"]
def add_number(trial: optuna.Trial, model_dir: Path):
@@ -109,12 +120,13 @@ def run_trial(
trial.number = number
# Add user attributes
trial._user_attrs.update(user_attrs)
[trial.set_user_attr(k, v) for k, v in user_attrs.items()]
print('trial', trial.number, trial, trial.params, trial.user_attrs)
model, trainer = main(
trial, PL_MODEL_CLS, name=name, MODEL_DIR=MODEL_DIR, train=False, prune=False
)
checkpoints = sorted(Path(trainer.checkpoint_callback.dirpath).glob("*.ckpt"))
if len(checkpoints)==0 or number is None:
try:
neural_processes/utils.py
+2
View File
@@ -6,6 +6,7 @@ import torch
import math
import torch
import optuna
from .logger import logger
def init_random_seed(seed):
@@ -91,6 +92,7 @@ def hparams_power(hparams):
if k.endswith("_power"):
k_new = k.replace("_power", "")
hparams[k_new] = int(2 ** hparams[k])
logger.debug('hparams %s', hparams)
return hparams
def log_prob_sigma(value, loc, log_scale):
smartmeters.ipynb
+188 -908
View File
File diff suppressed because one or more lines are too long