From f40ac35cc87f01cdc1c6d41f8d3499fc7c945412 Mon Sep 17 00:00:00 2001 From: Krystyna Milian Date: Sun, 14 Jul 2024 10:32:37 +0200 Subject: [PATCH] fixed PhaseNet normalization, added reducing lr on plateau --- config.json | 6 ++--- experiments/sweep_basicphase_ae.yaml | 11 +++----- experiments/sweep_eqtransformer.yaml | 8 ++---- experiments/sweep_gpd.yaml | 10 +++----- scripts/hyperparameter_sweep.py | 30 ++++++++++------------ scripts/input_validate.py | 18 ++++++++----- scripts/models.py | 50 +++++++++++++++++++++++++----------- 7 files changed, 72 insertions(+), 61 deletions(-) diff --git a/config.json b/config.json index acd11b9..b320b67 100644 --- a/config.json +++ b/config.json @@ -1,6 +1,6 @@ { - "dataset_name": "bogdanka", - "data_path": "datasets/bogdanka/seisbench_format/", + "dataset_name": "bogdanka_2018_2022", + "data_path": "datasets/bogdanka_2018_2022/seisbench_format/", "targets_path": "datasets/targets", "models_path": "weights", "configs_path": "experiments", @@ -13,5 +13,5 @@ "BasicPhaseAE": "sweep_basicphase_ae.yaml", "EQTransformer": "sweep_eqtransformer.yaml" }, - "experiment_count": 20 + "experiment_count": 15 } \ No newline at end of file diff --git a/experiments/sweep_basicphase_ae.yaml b/experiments/sweep_basicphase_ae.yaml index ba45c47..7e3bda1 100644 --- a/experiments/sweep_basicphase_ae.yaml +++ b/experiments/sweep_basicphase_ae.yaml @@ -1,3 +1,4 @@ +name: BasicPhaseAE method: bayes metric: goal: minimize @@ -7,13 +8,9 @@ parameters: value: - BasicPhaseAE batch_size: - distribution: int_uniform - max: 1024 - min: 256 + values: [64, 128, 256] max_epochs: value: - - 20 + - 30 learning_rate: - distribution: uniform - max: 0.02 - min: 0.001 + values: [0.01, 0.005, 0.001] diff --git a/experiments/sweep_eqtransformer.yaml b/experiments/sweep_eqtransformer.yaml index 21735f4..03741a2 100644 --- a/experiments/sweep_eqtransformer.yaml +++ b/experiments/sweep_eqtransformer.yaml @@ -8,13 +8,9 @@ parameters: value: - EQTransformer batch_size: - distribution: int_uniform - max: 1024 - min: 256 + values: [64, 128, 256] max_epochs: value: - 30 learning_rate: - distribution: uniform - max: 0.02 - min: 0.005 + values: [0.01, 0.005, 0.001] \ No newline at end of file diff --git a/experiments/sweep_gpd.yaml b/experiments/sweep_gpd.yaml index 58f3674..29db5bf 100644 --- a/experiments/sweep_gpd.yaml +++ b/experiments/sweep_gpd.yaml @@ -1,4 +1,4 @@ -name: GPD_fixed_highpass:2-10 +name: GPD method: bayes metric: goal: minimize @@ -8,16 +8,12 @@ parameters: value: - GPD batch_size: - distribution: int_uniform - max: 1024 - min: 256 + values: [64, 128, 256] max_epochs: value: - 30 learning_rate: - distribution: uniform - max: 0.02 - min: 0.005 + values: [0.01, 0.005, 0.001] highpass: value: - 1 diff --git a/scripts/hyperparameter_sweep.py b/scripts/hyperparameter_sweep.py index 981c38d..27ef805 100644 --- a/scripts/hyperparameter_sweep.py +++ b/scripts/hyperparameter_sweep.py @@ -56,6 +56,10 @@ def get_trainer_args(config): return trainer_args +def get_arg(arg): + if type(arg) == list: + return arg[0] + return arg class HyperparameterSweep: def __init__(self, project_name, sweep_config): @@ -87,10 +91,10 @@ class HyperparameterSweep: return all_not_running def run_experiment(self): - try: logger.info("Starting a new run...") + run = wandb.init( project=self.project_name, config=config_loader.config, @@ -103,30 +107,24 @@ class HyperparameterSweep: exclude_fn=lambda path: path.endswith("template.sh") ) - model_name = wandb.config.model_name[0] + model_name = get_arg(wandb.config.model_name) model_args = models.get_model_specific_args(wandb.config) + if "pretrained" in wandb.config: - weights = wandb.config.get("pretrained") - if type(weights) == list: - weights = weights[0] + weights = get_arg(wandb.config.pretrained) if weights != "false": model_args["pretrained"] = weights - if "norm" in wandb.config: - model_args["norm"] = wandb.config.norm - logger.debug(f"Initializing {model_name}") + if "norm" in wandb.config: + model_args["norm"] = get_arg(wandb.config.norm) + if "finetuning" in wandb.config: - # train for a few epochs with some frozen params, then unfreeze and continue training - if type(wandb.config.finetuning) == list: - finetuning_strategy = wandb.config.finetuning[0] - else: - finetuning_strategy = wandb.config.finetuning - model_args['finetuning_strategy'] = finetuning_strategy + model_args['finetuning_strategy'] = get_arg(wandb.config.finetuning) if "lr_reduce_factor" in wandb.config: - model_args['steplr_gamma'] = wandb.config.lr_reduce_factor - + model_args['steplr_gamma'] = get_arg(wandb.config.lr_reduce_factor) + logger.debug(f"Initializing {model_name} with args: {model_args}") model = models.__getattribute__(model_name + "Lit")(**model_args) train_loader, dev_loader = train.prepare_data(wandb.config, model, test_run=False) diff --git a/scripts/input_validate.py b/scripts/input_validate.py index fa6a8ef..c6bdfa4 100644 --- a/scripts/input_validate.py +++ b/scripts/input_validate.py @@ -85,6 +85,9 @@ class PhaseNetParameters(Parameters): finetuning: Finetuning = None lr_reduce_factor: Optional[Union[FloatDistribution, NumericValue, NumericValues]] = None + highpass: Union[NumericValue, NumericValues, FloatDistribution, IntDistribution] = None + lowpass: Union[NumericValue, NumericValues, FloatDistribution, IntDistribution] = None + @field_validator("model_name") def validate_model(cls, v): if "PhaseNet" not in v.value: @@ -92,23 +95,24 @@ class PhaseNetParameters(Parameters): return v -class GPDParameters(Parameters): +class FilteringParameters(Parameters): model_config = ConfigDict(extra='forbid') - highpass: Union[NumericValue, NumericValues, FloatDistribution] = None - lowpass: Union[NumericValue, NumericValues, FloatDistribution] = None + highpass: Union[NumericValue, NumericValues, FloatDistribution, IntDistribution] = None + lowpass: Union[NumericValue, NumericValues, FloatDistribution, IntDistribution] = None @field_validator("model_name") def validate_model(cls, v): - if "GPD" not in v.value: - raise ValueError("Additional parameters implemented for GPD only") + print(v.value) + if v.value[0] not in ["GPD", "PhaseNet"]: + raise ValueError("Filtering parameters implemented for GPD and PhaseNet only") class InputParams(BaseModel): name: str method: str metric: Metric - parameters: Union[Parameters, PhaseNetParameters, GPDParameters] + parameters: Union[Parameters, PhaseNetParameters, FilteringParameters] def validate_sweep_yaml(yaml_filename, model_name=None): @@ -138,5 +142,5 @@ def validate_sweep_config(sweep_config, model_name=None): if __name__ == "__main__": - yaml_filename = "../experiments/sweep_phasenet_lumineos_lr_bs.yaml" + yaml_filename = "../experiments/sweep_phasenet_bogdanka_lr_bs.yaml" validate_sweep_yaml(yaml_filename, None) diff --git a/scripts/models.py b/scripts/models.py index f35b0d9..7956c44 100644 --- a/scripts/models.py +++ b/scripts/models.py @@ -143,6 +143,9 @@ class PhaseNetLit(SeisBenchModuleLit): self.loss = vector_cross_entropy self.pretrained = kwargs.pop("pretrained", None) self.norm = kwargs.pop("norm", "peak") + self.highpass = kwargs.pop("highpass", None) + self.lowpass = kwargs.pop("lowpass", None) + if self.pretrained is not None: self.model = sbm.PhaseNet.from_pretrained(self.pretrained) @@ -152,6 +155,7 @@ class PhaseNetLit(SeisBenchModuleLit): self.finetuning_strategy = kwargs.pop("finetuning_strategy", None) self.steplr_gamma = kwargs.pop("steplr_gamma", 0.1) + self.reduce_lr_on_plateau = False self.initial_epochs = 0 @@ -163,36 +167,33 @@ class PhaseNetLit(SeisBenchModuleLit): self.freeze() - def forward(self, x): return self.model(x) - def shared_step(self, batch): x = batch["X"] y_true = batch["y"] y_pred = self.model(x) return self.loss(y_pred, y_true) - def training_step(self, batch, batch_idx): loss = self.shared_step(batch) self.log("train_loss", loss) return loss - def validation_step(self, batch, batch_idx): loss = self.shared_step(batch) self.log("val_loss", loss) return loss - def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.lr) if self.finetuning_strategy is not None: scheduler = lr_scheduler.LambdaLR(optimizer, self.lr_lambda) + self.reduce_lr_on_plateau = False else: scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=3) + self.reduce_lr_on_plateau = True # return { 'optimizer': optimizer, @@ -200,11 +201,10 @@ class PhaseNetLit(SeisBenchModuleLit): 'scheduler': scheduler, 'monitor': 'val_loss', 'interval': 'epoch', - 'reduce_on_plateau': False, + 'reduce_on_plateau': self.reduce_lr_on_plateau, }, } - def lr_lambda(self, epoch): # reduce lr after x initial epochs if epoch == self.initial_epochs: @@ -212,15 +212,25 @@ class PhaseNetLit(SeisBenchModuleLit): return self.lr - def lr_scheduler_step(self, scheduler, metric): - scheduler.step(epoch=self.current_epoch) - + if self.reduce_lr_on_plateau: + scheduler.step(metric, epoch=self.current_epoch) + else: + scheduler.step(epoch=self.current_epoch) # def lr_scheduler_step(self, scheduler, optimizer_idx, metric): # scheduler.step(epoch=self.current_epoch) def get_augmentations(self): + filter = [] + if self.highpass is not None: + filter = [sbg.Filter(1, self.highpass, "highpass", forward_backward=True)] + logger.info(f"Using highpass filer {self.highpass}") + if self.lowpass is not None: + filter += [sbg.Filter(1, self.lowpass, "lowpass", forward_backward=True)] + logger.info(f"Using lowpass filer {self.lowpass}") + logger.info(filter) + return [ # In 2/3 of the cases, select windows around picks, to reduce amount of noise traces in training. # Uses strategy variable, as padding will be handled by the random window. @@ -244,22 +254,28 @@ class PhaseNetLit(SeisBenchModuleLit): windowlen=3001, strategy="pad", ), - sbg.ChangeDtype(np.float32), sbg.Normalize(demean_axis=-1, amp_norm_axis=-1, amp_norm_type=self.norm), + *filter, + sbg.ChangeDtype(np.float32), sbg.ProbabilisticLabeller( label_columns=phase_dict, sigma=self.sigma, dim=0 ), ] - def get_eval_augmentations(self): + filter = [] + if self.highpass is not None: + filter = [sbg.Filter(1, self.highpass, "highpass", forward_backward=True)] + if self.lowpass is not None: + filter += [sbg.Filter(1, self.lowpass, "lowpass", forward_backward=True)] + return [ sbg.SteeredWindow(windowlen=3001, strategy="pad"), - sbg.ChangeDtype(np.float32), sbg.Normalize(demean_axis=-1, amp_norm_axis=-1, amp_norm_type=self.norm), + *filter, + sbg.ChangeDtype(np.float32), ] - def predict_step(self, batch, batch_idx=None, dataloader_idx=None): x = batch["X"] window_borders = batch["window_borders"] @@ -1211,8 +1227,12 @@ def get_model_specific_args(config): if 'highpass' in config: args['highpass'] = config.highpass if 'lowpass' in config: - args['lowpass'] = config.lowpass[0] + args['lowpass'] = config.lowpass case 'PhaseNet': + if 'highpass' in config: + args['highpass'] = config.highpass + if 'lowpass' in config: + args['lowpass'] = config.lowpass if 'sample_boundaries' in config: args['sample_boundaries'] = config.sample_boundaries case 'DPPPicker':