From 27bf5f8922250823616729a4987e1c97ae702f0e Mon Sep 17 00:00:00 2001 From: xiaoyewww <641311428@qq.com> Date: Thu, 19 Dec 2024 23:51:34 +0800 Subject: [PATCH] feat(neuralop): add the rest of tests --- neuralop/models/tests/__init__.py | 0 .../models/tests}/test_fno.py | 0 .../models/tests}/test_fnogno.py | 0 .../models/tests}/test_uno.py | 0 neuralop/tests/__init__.py | 0 neuralop/tests/test_config.yaml | 57 +++++ neuralop/tests/test_config_key.txt | 1 + neuralop/tests/test_model_from_config.py | 44 ++++ neuralop/tests/test_utils.py | 103 +++++++++ neuralop/training/tests/test_callbacks.py | 203 ++++++++++++++++++ 10 files changed, 408 insertions(+) create mode 100644 neuralop/models/tests/__init__.py rename {examples/models => neuralop/models/tests}/test_fno.py (100%) rename {examples/models => neuralop/models/tests}/test_fnogno.py (100%) rename {examples/models => neuralop/models/tests}/test_uno.py (100%) create mode 100644 neuralop/tests/__init__.py create mode 100644 neuralop/tests/test_config.yaml create mode 100644 neuralop/tests/test_config_key.txt create mode 100644 neuralop/tests/test_model_from_config.py create mode 100644 neuralop/tests/test_utils.py create mode 100644 neuralop/training/tests/test_callbacks.py diff --git a/neuralop/models/tests/__init__.py b/neuralop/models/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/models/test_fno.py b/neuralop/models/tests/test_fno.py similarity index 100% rename from examples/models/test_fno.py rename to neuralop/models/tests/test_fno.py diff --git a/examples/models/test_fnogno.py b/neuralop/models/tests/test_fnogno.py similarity index 100% rename from examples/models/test_fnogno.py rename to neuralop/models/tests/test_fnogno.py diff --git a/examples/models/test_uno.py b/neuralop/models/tests/test_uno.py similarity index 100% rename from examples/models/test_uno.py rename to neuralop/models/tests/test_uno.py diff --git a/neuralop/tests/__init__.py b/neuralop/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/neuralop/tests/test_config.yaml b/neuralop/tests/test_config.yaml new file mode 100644 index 0000000..7cbaa05 --- /dev/null +++ b/neuralop/tests/test_config.yaml @@ -0,0 +1,57 @@ +default: &DEFAULT + + #General + verbose: True + arch: 'tfno2d' + + # FNO related + tfno2d: + data_channels: 3 + n_modes_height: 8 + n_modes_width: 8 + hidden_channels: 32 + projection_channels: 32 + n_layers: 2 + domain_padding: 0 + domain_padding_mode: 'symmetric' + fft_norm: 'forward' + norm: None + skip: 'soft-gating' + implementation: 'factorized' + + use_mlp: 1 + mlp: + expansion: 0.5 + dropout: 0 + + factorization: None + rank: 1.0 + fixed_rank_modes: None + dropout: 0.0 + tensor_lasso_penalty: 0.0 + joint_factorization: False + + data: + batch_size: 4 + n_train: 10 + size: 32 + + # Optimizer + opt: + n_epochs: 500 + learning_rate: 1e-3 + training_loss: 'h1' + weight_decay: 1e-4 + amp_autocast: True + + scheduler_T_max: 500 # For cosine only, typically take n_epochs + scheduler_patience: 5 # For ReduceLROnPlateau only + scheduler: 'StepLR' # Or 'CosineAnnealingLR' OR 'ReduceLROnPlateau' + step_size: 100 + gamma: 0.5 + + # Patching + patching: + levels: 0 + padding: 0 #.1 + stitching: True diff --git a/neuralop/tests/test_config_key.txt b/neuralop/tests/test_config_key.txt new file mode 100644 index 0000000..d6d6f03 --- /dev/null +++ b/neuralop/tests/test_config_key.txt @@ -0,0 +1 @@ +my_secret_key \ No newline at end of file diff --git a/neuralop/tests/test_model_from_config.py b/neuralop/tests/test_model_from_config.py new file mode 100644 index 0000000..57d4e75 --- /dev/null +++ b/neuralop/tests/test_model_from_config.py @@ -0,0 +1,44 @@ + +import paddle +import time +from tensorly import tenalg +tenalg.set_backend('einsum') +from pathlib import Path + +from configmypy import ConfigPipeline, YamlConfig +from neuralop import get_model + +def test_from_config(): + """Test forward/backward from a config file""" + # Read the configuration + config_name = 'default' + config_path = Path(__file__).parent.as_posix() + pipe = ConfigPipeline([YamlConfig('./test_config.yaml', config_name=config_name, config_folder=config_path), + ]) + config = pipe.read_conf() + config_name = pipe.steps[-1].config_name + + batch_size = config.data.batch_size + size = config.data.size + + if paddle.device.cuda.device_count() >= 1: + device = 'cuda' + else: + device = 'cpu' + + paddle.device.set_device(device=device) + + model = get_model(config) + model = model + + in_data = paddle.randn([batch_size, 3, size, size]) + print(model.__class__) + print(model) + + t1 = time.time() + out = model(in_data) + t = time.time() - t1 + print(f'Output of size {out.shape} in {t}.') + + loss = out.sum() + loss.backward() diff --git a/neuralop/tests/test_utils.py b/neuralop/tests/test_utils.py new file mode 100644 index 0000000..e92d32a --- /dev/null +++ b/neuralop/tests/test_utils.py @@ -0,0 +1,103 @@ +from ..utils import get_wandb_api_key, wandb_login +from ..utils import count_model_params, count_tensor_params +from pathlib import Path +import pytest +import wandb +import os +import paddle +from paddle import nn + + +def test_count_model_params(): + # A nested dummy model to make sure all parameters are counted + class DumyModel(nn.Layer): + def __init__(self, n_submodels=0, dtype=paddle.float32): + super().__init__() + + self.n_submodels = n_submodels + self.param = paddle.base.framework.EagerParamBase.from_tensor(paddle.randn((2, 3, 4), dtype=dtype)) + if n_submodels: + self.model = DumyModel(n_submodels - 1, dtype=dtype) + + n_submodels = 2 + model = DumyModel(n_submodels=n_submodels) + n_params = count_model_params(model) + print(model) + assert n_params == (n_submodels+1) * 2 * 3 * 4 + + model = DumyModel(n_submodels=n_submodels, dtype=paddle.complex64) + n_params = count_model_params(model) + print(model) + assert n_params == 2 * (n_submodels+1) * 2*3*4 + + +def test_count_tensor_params(): + # Case 1 : real tensor + x = paddle.randn((2, 3, 4, 5, 6), dtype=paddle.float32) + + # dims = None: count all params + n_params = count_tensor_params(x) + assert n_params == 2*3*4*5*6 + # Only certain dims + n_params = count_tensor_params(x, dims=[1, 3]) + assert n_params == 3*5 + + # Case 2 : complex tensor + x = paddle.randn((2, 3, 4, 5, 6), dtype=paddle.complex64) + + # dims = None: count all params + n_params = count_tensor_params(x) + assert n_params == 2*3*4*5*6 * 2 + # Only certain dims + n_params = count_tensor_params(x, dims=[1, 3]) + assert n_params == 3*5 * 2 + + +def test_get_wandb_api_key(): + # Make sure no env var key set + os.environ.pop("WANDB_API_KEY", None) + + # Read from file + filepath = Path(__file__).parent.joinpath('test_config_key.txt').as_posix() + key = get_wandb_api_key(filepath) + assert key == 'my_secret_key' + + # Read from env var + os.environ["WANDB_API_KEY"] = 'key_from_env' + key = get_wandb_api_key(filepath) + assert key == 'key_from_env' + + # Read from env var with incorrect file + os.environ["WANDB_API_KEY"] = 'key_from_env' + key = get_wandb_api_key('wrong_path') + assert key == 'key_from_env' + + +def test_ArgparseConfig(monkeypatch): + def login(key): + if key == 'my_secret_key': + return True + + raise ValueError('Wrong key') + + monkeypatch.setattr(wandb, "login", login) + + # Make sure no env var key set + os.environ.pop("WANDB_API_KEY", None) + + # Read from file + filepath = Path(__file__).parent.joinpath('test_config_key.txt').as_posix() + assert wandb_login(filepath) is None + + # Read from env var + os.environ["WANDB_API_KEY"] = 'my_secret_key' + assert wandb_login() is None + + # Read from env var + os.environ["WANDB_API_KEY"] = 'wrong_key' + assert wandb_login(key='my_secret_key') is None + + # Read from env var + os.environ["WANDB_API_KEY"] = 'wrong_key' + with pytest.raises(ValueError): + wandb_login() diff --git a/neuralop/training/tests/test_callbacks.py b/neuralop/training/tests/test_callbacks.py new file mode 100644 index 0000000..5b8f4b4 --- /dev/null +++ b/neuralop/training/tests/test_callbacks.py @@ -0,0 +1,203 @@ +import shutil +from pathlib import Path + +import paddle +from paddle import nn +from paddle.io import Dataset, DataLoader + +from neuralop import Trainer, LpLoss, H1Loss, CheckpointCallback +from neuralop.models.base_model import BaseModel + + +class DummyDataset(Dataset): + # Simple linear regression problem, PyTorch style + + def __init__(self, n_examples): + super().__init__() + + self.X = paddle.randn((n_examples, 50)) + self.y = paddle.randn((n_examples, 1)) + + def __getitem__(self, idx): + return {'x': self.X[idx], 'y': self.y[idx]} + + def __len__(self): + return self.X.shape[0] + + +class DummyModel(BaseModel, name='Dummy'): + """ + Simple linear model to mock-up our model API + """ + + def __init__(self, features, **kwargs): + super().__init__() + self.net = nn.Linear(features, 1) + + def forward(self, x, **kwargs): + """ + Throw out extra args as in FNO and other models + """ + return self.net(x) + + +def test_model_checkpoint_saves(): + save_pth = Path('./test_checkpoints') + + model = DummyModel(50) + + train_loader = DataLoader(DummyDataset(100)) + + trainer = Trainer(model=model, + n_epochs=5, + callbacks=[ + CheckpointCallback(save_dir=save_pth, + save_optimizer=True, + save_scheduler=True) + ]) + + scheduler = paddle.optimizer.lr.CosineAnnealingDecay(learning_rate=8e-3, T_max=30) + optimizer = paddle.optimizer.Adam( + learning_rate=scheduler, parameters=model.parameters(), weight_decay=1e-4 + ) + + # Creating the losses + l2loss = LpLoss(d=2, p=2) + + trainer.train(train_loader=train_loader, + test_loaders={}, + optimizer=optimizer, + scheduler=scheduler, + regularizer=None, + training_loss=l2loss, + eval_losses=None, + ) + + for file_ext in ['model_state_dict.pdmodel', 'model_metadata.pkl', 'optimizer.pdopt', 'scheduler.pdopt']: + file_pth = save_pth / file_ext + assert file_pth.exists() + + # clean up dummy checkpoint directory after testing + shutil.rmtree('./test_checkpoints') + + +def test_model_checkpoint_and_resume(): + save_pth = Path('./full_states') + model = DummyModel(50) + + train_loader = DataLoader(DummyDataset(100)) + test_loader = DataLoader(DummyDataset(20)) + + trainer = Trainer( + model=model, + n_epochs=5, + callbacks=[ + CheckpointCallback( + save_dir=save_pth, save_optimizer=True, save_scheduler=True, save_best='h1' + ) # monitor h1 loss + ] + ) + + scheduler = paddle.optimizer.lr.CosineAnnealingDecay(learning_rate=8e-3, T_max=30) + optimizer = paddle.optimizer.Adam( + learning_rate=scheduler, parameters=model.parameters(), weight_decay=1e-4 + ) + + # Creating the losses + l2loss = LpLoss(d=2, p=2) + h1loss = H1Loss(d=2) + + eval_losses = {'h1': h1loss, 'l2': l2loss} + + trainer.train(train_loader=train_loader, + test_loaders={'': test_loader}, + optimizer=optimizer, + scheduler=scheduler, + regularizer=None, + training_loss=l2loss, + eval_losses=eval_losses + ) + + for file_ext in ['best_model_state_dict.pdmodel', 'best_model_metadata.pkl', 'optimizer.pdopt', 'scheduler.pdopt']: + file_pth = save_pth / file_ext + assert file_pth.exists() + + # Resume from checkpoint + trainer = Trainer( + model=model, + n_epochs=5, + callbacks=[ + CheckpointCallback(save_dir='./checkpoints', resume_from_dir='./full_states') + ] + ) + + trainer.train( + train_loader=train_loader, + test_loaders={'': test_loader}, + optimizer=optimizer, + scheduler=scheduler, + regularizer=None, + training_loss=l2loss, + eval_losses=eval_losses, + ) + + # clean up dummy checkpoint directory after testing + shutil.rmtree(save_pth) + + +# ensure that model accuracy after loading from checkpoint +# is comparable to accuracy at time of save +def test_load_from_checkpoint(): + model = DummyModel(50) + + train_loader = DataLoader(DummyDataset(100)) + test_loader = DataLoader(DummyDataset(20)) + + trainer = Trainer( + model=model, + n_epochs=5, + callbacks=[ + CheckpointCallback( + save_dir='./full_states', + save_optimizer=True, + save_scheduler=True, + save_best='h1' + ) # monitor h1 loss + ] + ) + + scheduler = paddle.optimizer.lr.CosineAnnealingDecay(learning_rate=8e-3, T_max=30) + optimizer = paddle.optimizer.Adam( + learning_rate=scheduler, parameters=model.parameters(), weight_decay=1e-4 + ) + + # Creating the losses + l2loss = LpLoss(d=2, p=2) + h1loss = H1Loss(d=2) + + eval_losses = {'h1': h1loss, 'l2': l2loss} + + orig_model_eval_errors = trainer.train( + train_loader=train_loader, + test_loaders={'': test_loader}, + optimizer=optimizer, + scheduler=scheduler, + regularizer=None, + training_loss=l2loss, + eval_losses=eval_losses + ) + + # create a new model from saved checkpoint and evaluate + loaded_model = DummyModel.from_checkpoint(save_folder='./full_states', save_name='best_model') + trainer = Trainer( + model=loaded_model, + n_epochs=1, + ) + + loaded_model_eval_errors = trainer.evaluate(loss_dict=eval_losses, data_loader=test_loader) + + # log prefix is empty except for default underscore + assert orig_model_eval_errors['_h1'] - loaded_model_eval_errors['_h1'] < 0.1 + + # clean up dummy checkpoint directory after testing + shutil.rmtree('./full_states')