diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 0000000..7238934 --- /dev/null +++ b/.coveragerc @@ -0,0 +1,4 @@ +[run] +omit = + # omit anything in the contrib directory + */contrib/* \ No newline at end of file diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ceab823..7e1d97d 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -3,6 +3,8 @@ Unreleased_ Added ^^^^^ +* `@senwu`_: Support more unit tests. +* `@senwu`_: Support all pytorch optimizers. * `@senwu`_: Support accuracy@k metric. * `@senwu`_: Support cosine annealing lr scheduler. diff --git a/docs/user/config.rst b/docs/user/config.rst index 64a52ea..63c49be 100644 --- a/docs/user/config.rst +++ b/docs/user/config.rst @@ -38,26 +38,61 @@ The default ``.emmental-config.yaml`` configuration file is shown below:: train_split: train # the split for training, accepts str or list of strs valid_split: valid # the split for validation, accepts str or list of strs test_split: test # the split for testing, accepts str or list of strs - ignore_index: 0 # the ignore index, uses for masking samples + ignore_index: # the ignore index, uses for masking samples optimizer_config: optimizer: adam # [sgd, adam, adamax, bert_adam] lr: 0.001 # Learing rate l2: 0.0 # l2 regularization - grad_clip: 1.0 # gradient clipping - sgd_config: - momentum: 0.9 + grad_clip: # gradient clipping + asgd_config: + lambd: 0.0001 + alpha: 0.75 + t0: 1000000.0 + adadelta_config: + rho: 0.9 + eps: 0.000001 + adagrad_config: + lr_decay: 0 + initial_accumulator_value: 0 + eps: 0.0000000001 adam_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 amsgrad: False + adamw_config: + betas: !!python/tuple [0.9, 0.999] + eps: 0.00000001 + amsgrad: False adamax_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 + lbfgs_config: + max_iter: 20 + max_eval: + tolerance_grad: 0.0000001 + tolerance_change: 0.000000001 + history_size: 100 + line_search_fn: + rms_prop_config: + alpha: 0.99 + eps: 0.00000001 + momentum: 0 + centered: False + r_prop_config: + etas: !!python/tuple [0.5, 1.2] + step_sizes: !!python/tuple [0.000001, 50] + sgd_config: + momentum: 0 + dampening: 0 + nesterov: False + sparse_adam_config: + betas: !!python/tuple [0.9, 0.999] + eps: 0.00000001 bert_adam_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 lr_scheduler_config: - lr_scheduler: # [linear, exponential, reduce_on_plateau] + lr_scheduler: # [linear, exponential, reduce_on_plateau, cosine_annealing] warmup_steps: # warm up steps warmup_unit: batch # [epoch, batch] warmup_percentage: # warm up percentage @@ -79,6 +114,8 @@ The default ``.emmental-config.yaml`` configuration file is shown below:: - 1000 gamma: 0.1 last_epoch: -1 + cosine_annealing_config: + last_epoch: -1 task_scheduler_config: task_scheduler: round_robin # [sequential, round_robin, mixed] sequential_scheduler_config: @@ -92,7 +129,7 @@ The default ``.emmental-config.yaml`` configuration file is shown below:: # Logging configuration logging_config: counter_unit: epoch # [epoch, batch] - evaluation_freq: 2 + evaluation_freq: 1 writer_config: writer: tensorboard # [json, tensorboard] verbose: True diff --git a/docs/user/learning.rst b/docs/user/learning.rst index 64825e8..8753d65 100644 --- a/docs/user/learning.rst +++ b/docs/user/learning.rst @@ -28,26 +28,61 @@ The learning parameters of the model are described below:: train_split: train # the split for training, accepts str or list of strs valid_split: valid # the split for validation, accepts str or list of strs test_split: test # the split for testing, accepts str or list of strs - ignore_index: 0 # the ignore index, uses for masking samples + ignore_index: # the ignore index, uses for masking samples optimizer_config: optimizer: adam # [sgd, adam, adamax, bert_adam] lr: 0.001 # Learing rate l2: 0.0 # l2 regularization - grad_clip: 1.0 # gradient clipping - sgd_config: - momentum: 0.9 + grad_clip: # gradient clipping + asgd_config: + lambd: 0.0001 + alpha: 0.75 + t0: 1000000.0 + adadelta_config: + rho: 0.9 + eps: 0.000001 + adagrad_config: + lr_decay: 0 + initial_accumulator_value: 0 + eps: 0.0000000001 adam_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 amsgrad: False + adamw_config: + betas: !!python/tuple [0.9, 0.999] + eps: 0.00000001 + amsgrad: False adamax_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 + lbfgs_config: + max_iter: 20 + max_eval: + tolerance_grad: 0.0000001 + tolerance_change: 0.000000001 + history_size: 100 + line_search_fn: + rms_prop_config: + alpha: 0.99 + eps: 0.00000001 + momentum: 0 + centered: False + r_prop_config: + etas: !!python/tuple [0.5, 1.2] + step_sizes: !!python/tuple [0.000001, 50] + sgd_config: + momentum: 0 + dampening: 0 + nesterov: False + sparse_adam_config: + betas: !!python/tuple [0.9, 0.999] + eps: 0.00000001 bert_adam_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 lr_scheduler_config: - lr_scheduler: # [linear, exponential, reduce_on_plateau] + lr_scheduler: # [linear, exponential, reduce_on_plateau, cosine_annealing] warmup_steps: # warm up steps warmup_unit: batch # [epoch, batch] warmup_percentage: # warm up percentage @@ -69,6 +104,8 @@ The learning parameters of the model are described below:: - 1000 gamma: 0.1 last_epoch: -1 + cosine_annealing_config: + last_epoch: -1 task_scheduler_config: task_scheduler: round_robin # [sequential, round_robin, mixed] sequential_scheduler_config: diff --git a/docs/user/logging.rst b/docs/user/logging.rst index 2652c1c..56612d4 100644 --- a/docs/user/logging.rst +++ b/docs/user/logging.rst @@ -25,7 +25,7 @@ The logging parameters of Emmental are described below:: # Logging configuration logging_config: counter_unit: epoch # [epoch, batch] - evaluation_freq: 2 + evaluation_freq: 1 writer_config: writer: tensorboard # [json, tensorboard] verbose: True diff --git a/src/emmental/emmental-default-config.yaml b/src/emmental/emmental-default-config.yaml index 74770a3..855af40 100644 --- a/src/emmental/emmental-default-config.yaml +++ b/src/emmental/emmental-default-config.yaml @@ -28,17 +28,50 @@ learner_config: lr: 0.001 # Learing rate l2: 0.0 # l2 regularization grad_clip: # gradient clipping - sgd_config: - momentum: 0.9 - dampening: 0 - nesterov: False + asgd_config: + lambd: 0.0001 + alpha: 0.75 + t0: 1000000.0 + adadelta_config: + rho: 0.9 + eps: 0.000001 + adagrad_config: + lr_decay: 0 + initial_accumulator_value: 0 + eps: 0.0000000001 adam_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 amsgrad: False + adamw_config: + betas: !!python/tuple [0.9, 0.999] + eps: 0.00000001 + amsgrad: False adamax_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 + lbfgs_config: + max_iter: 20 + max_eval: + tolerance_grad: 0.0000001 + tolerance_change: 0.000000001 + history_size: 100 + line_search_fn: + rms_prop_config: + alpha: 0.99 + eps: 0.00000001 + momentum: 0 + centered: False + r_prop_config: + etas: !!python/tuple [0.5, 1.2] + step_sizes: !!python/tuple [0.000001, 50] + sgd_config: + momentum: 0 + dampening: 0 + nesterov: False + sparse_adam_config: + betas: !!python/tuple [0.9, 0.999] + eps: 0.00000001 bert_adam_config: betas: !!python/tuple [0.9, 0.999] eps: 0.00000001 diff --git a/src/emmental/learner.py b/src/emmental/learner.py index 95d53a3..b982365 100644 --- a/src/emmental/learner.py +++ b/src/emmental/learner.py @@ -52,41 +52,43 @@ def _set_optimizer(self, model: EmmentalModel) -> None: model(EmmentalModel): The model to set up the optimizer. """ - - # TODO: add more optimizer support and fp16 optimizer_config = Meta.config["learner_config"]["optimizer_config"] opt = optimizer_config["optimizer"] parameters = filter(lambda p: p.requires_grad, model.parameters()) - if opt == "sgd": - optimizer = optim.SGD( - parameters, - lr=optimizer_config["lr"], - weight_decay=optimizer_config["l2"], - **optimizer_config["sgd_config"], - ) - elif opt == "adam": - optimizer = optim.Adam( + optim_dict = { + # PyTorch optimizer + "asgd": optim.ASGD, + "adadelta": optim.Adadelta, + "adagrad": optim.Adagrad, + "adam": optim.Adam, + "adamw": optim.AdamW, + "adamax": optim.Adamax, + "lbfgs": optim.LBFGS, + "rms_prop": optim.RMSprop, + "r_prop": optim.Rprop, + "sgd": optim.SGD, + "sparse_adam": optim.SparseAdam, + # Customize optimizer + "bert_adam": BertAdam, + } + + if opt in ["lbfgs", "r_prop", "sparse_adam"]: + optimizer = optim_dict[opt]( parameters, lr=optimizer_config["lr"], - weight_decay=optimizer_config["l2"], - **optimizer_config["adam_config"], - ) - elif opt == "adamax": - optimizer = optim.Adamax( - parameters, - lr=optimizer_config["lr"], - weight_decay=optimizer_config["l2"], - **optimizer_config["adamax_config"], + **optimizer_config[f"{opt}_config"], ) - elif opt == "bert_adam": - optimizer = BertAdam( + elif opt in optim_dict.keys(): + optimizer = optim_dict[opt]( parameters, lr=optimizer_config["lr"], weight_decay=optimizer_config["l2"], - **optimizer_config["bert_adam_config"], + **optimizer_config[f"{opt}_config"], ) + elif isinstance(opt, optim.Optimizer): + optimizer = opt(parameters) else: raise ValueError(f"Unrecognized optimizer option '{opt}'") diff --git a/src/emmental/modules/rnn_module.py b/src/emmental/modules/rnn_module.py index 210139a..83ad822 100644 --- a/src/emmental/modules/rnn_module.py +++ b/src/emmental/modules/rnn_module.py @@ -105,12 +105,15 @@ def forward(self, x: Tensor, x_mask: Optional[Tensor] = None) -> Tensor: """ Mean pooling """ + if x_mask is None: + x_mask = x.new_full(x.size()[:2], fill_value=0, dtype=torch.uint8) x_lens = x_mask.data.eq(0).long().sum(dim=1) - weights = torch.ones(x.size()) / x_lens.unsqueeze(1).float() + weights = ( + output_word.new_ones(output_word.size()) + / x_lens.view(x_lens.size()[0], 1, 1).float() + ) weights.data.masked_fill_(x_mask.data.unsqueeze(dim=2), 0.0) - word_vectors = torch.bmm( - output_word.transpose(1, 2), weights.unsqueeze(2) - ).squeeze(2) + word_vectors = torch.sum(output_word * weights, dim=1) output = self.linear(word_vectors) if self.final_linear else word_vectors return output diff --git a/src/emmental/utils/parse_arg.py b/src/emmental/utils/parse_args.py similarity index 66% rename from src/emmental/utils/parse_arg.py rename to src/emmental/utils/parse_args.py index f8ebbf4..3cf88a8 100644 --- a/src/emmental/utils/parse_arg.py +++ b/src/emmental/utils/parse_args.py @@ -11,7 +11,7 @@ ) -def parse_arg(parser: Optional[ArgumentParser] = None) -> ArgumentParser: +def parse_args(parser: Optional[ArgumentParser] = None) -> ArgumentParser: r"""Parse the configuration from command line. Args: @@ -137,9 +137,23 @@ def parse_arg(parser: Optional[ArgumentParser] = None) -> ArgumentParser: optimizer_config.add_argument( "--optimizer", - type=str, + type=nullable_string, default="adam", - choices=["adam", "adamax", "sgd", "bert_adam"], + choices=[ + "asgd", + "adadelta", + "adagrad", + "adam", + "adamw", + "adamax", + "lbfgs", + "rms_prop", + "r_prop", + "sgd", + "sparse_adam", + "bert_adam", + None, + ], help="The optimizer to use", ) @@ -153,29 +167,190 @@ def parse_arg(parser: Optional[ArgumentParser] = None) -> ArgumentParser: "--grad_clip", type=nullable_float, default=None, help="Gradient clipping" ) + # ASGD config optimizer_config.add_argument( - "--sgd_momentum", type=float, default=0.9, help="SGD momentum" + "--asgd_lambd", type=float, default=0.0001, help="ASGD lambd" ) optimizer_config.add_argument( - "--sgd_dampening", type=float, default=0, help="SGD dampening" + "--asgd_alpha", type=float, default=0.75, help="ASGD alpha" ) optimizer_config.add_argument( - "--sgd_nesterov", type=str2bool, default=False, help="SGD nesterov" + "--asgd_t0", type=float, default=1000000.0, help="ASGD t0" + ) + + # Adadelta config + optimizer_config.add_argument( + "--adadelta_rho", type=float, default=0.9, help="Adadelta rho" + ) + + optimizer_config.add_argument( + "--adadelta_eps", type=float, default=0.000001, help="Adadelta eps" ) - # TODO: add adam/adamax/bert_adam betas + # Adagrad config + optimizer_config.add_argument( + "--adagrad_lr_decay", type=float, default=0, help="Adagrad lr_decay" + ) optimizer_config.add_argument( - "--amsgrad", + "--adagrad_initial_accumulator_value", + type=float, + default=0, + help="Adagrad initial accumulator value", + ) + + optimizer_config.add_argument( + "--adagrad_eps", type=float, default=0.0000000001, help="Adagrad eps" + ) + + # Adam config + optimizer_config.add_argument( + "--adam_betas", nargs="+", type=float, default=(0.9, 0.999), help="Adam betas" + ) + + optimizer_config.add_argument( + "--adam_eps", type=float, default=1e-8, help="Adam eps" + ) + + optimizer_config.add_argument( + "--adam_amsgrad", type=str2bool, default=False, help="Whether to use the AMSGrad variant of adam", ) + # AdamW config optimizer_config.add_argument( - "--eps", type=float, default=1e-8, help="eps in adam, adamax, or bert_adam" + "--adamw_betas", nargs="+", type=float, default=(0.9, 0.999), help="AdamW betas" + ) + + optimizer_config.add_argument( + "--adamw_eps", type=float, default=1e-8, help="AdamW eps" + ) + + optimizer_config.add_argument( + "--adamw_amsgrad", + type=str2bool, + default=False, + help="Whether to use the AMSGrad variant of AdamW", + ) + + # Adamax config + optimizer_config.add_argument( + "--adamax_betas", + nargs="+", + type=float, + default=(0.9, 0.999), + help="Adamax betas", + ) + + optimizer_config.add_argument( + "--adamax_eps", type=float, default=1e-8, help="Adamax eps" + ) + + # LBFGS config + optimizer_config.add_argument( + "--lbfgs_max_iter", type=int, default=20, help="LBFGS max iter" + ) + + optimizer_config.add_argument( + "--lbfgs_max_eval", type=nullable_int, default=None, help="LBFGS max eval" + ) + + optimizer_config.add_argument( + "--lbfgs_tolerance_grad", type=float, default=1e-07, help="LBFGS tolerance grad" + ) + + optimizer_config.add_argument( + "--lbfgs_tolerance_change", + type=float, + default=1e-09, + help="LBFGS tolerance change", + ) + + optimizer_config.add_argument( + "--lbfgs_history_size", type=int, default=100, help="LBFGS history size" + ) + + optimizer_config.add_argument( + "--lbfgs_line_search_fn", + type=nullable_string, + default=None, + help="LBFGS line search fn", + ) + + # RMSprop config + optimizer_config.add_argument( + "--rms_prop_alpha", type=float, default=0.99, help="RMSprop alpha" + ) + + optimizer_config.add_argument( + "--rms_prop_eps", type=float, default=1e-08, help="RMSprop eps" + ) + + optimizer_config.add_argument( + "--rms_prop_momentum", type=float, default=0, help="RMSprop momentum" + ) + + optimizer_config.add_argument( + "--rms_prop_centered", type=str2bool, default=False, help="RMSprop centered" + ) + + # Rprop config + optimizer_config.add_argument( + "--r_prop_etas", nargs="+", type=float, default=(0.5, 1.2), help="Rprop etas" + ) + + optimizer_config.add_argument( + "--r_prop_step_sizes", + nargs="+", + type=float, + default=(1e-06, 50), + help="Rprop step sizes", + ) + + # SGD config + optimizer_config.add_argument( + "--sgd_momentum", type=float, default=0.0, help="SGD momentum" + ) + + optimizer_config.add_argument( + "--sgd_dampening", type=float, default=0, help="SGD dampening" + ) + + optimizer_config.add_argument( + "--sgd_nesterov", type=str2bool, default=False, help="SGD nesterov" + ) + + # TODO: support etas and step_sizes in Rprop + # Rprop config + + # SparseAdam config + optimizer_config.add_argument( + "--sparse_adam_betas", + nargs="+", + type=float, + default=(0.9, 0.999), + help="SparseAdam betas", + ) + + optimizer_config.add_argument( + "--sparse_adam_eps", type=float, default=1e-08, help="SparseAdam eps" + ) + + # BertAdam config + optimizer_config.add_argument( + "--bert_adam_betas", + nargs="+", + type=float, + default=(0.9, 0.999), + help="BertAdam betas", + ) + + optimizer_config.add_argument( + "--bert_adam_eps", type=float, default=1e-08, help="BertAdam eps" ) # Scheduler configuration @@ -408,7 +583,7 @@ def parse_arg(parser: Optional[ArgumentParser] = None) -> ArgumentParser: return parser -def parse_arg_to_config(args: Namespace) -> Dict[str, Any]: +def parse_args_to_config(args: Namespace) -> Dict[str, Any]: r"""Parse the arguments to config dict Args: @@ -445,18 +620,59 @@ def parse_arg_to_config(args: Namespace) -> Dict[str, Any]: "lr": args.lr, "l2": args.l2, "grad_clip": args.grad_clip, + "asgd_config": { + "lambd": args.asgd_lambd, + "alpha": args.asgd_alpha, + "t0": args.asgd_t0, + }, + "adadelta_config": {"rho": args.adadelta_rho, "eps": args.adadelta_eps}, + "adagrad_config": { + "lr_decay": args.adagrad_lr_decay, + "initial_accumulator_value": args.adagrad_initial_accumulator_value, + "eps": args.adagrad_eps, + }, + "adam_config": { + "betas": args.adam_betas, + "amsgrad": args.adam_amsgrad, + "eps": args.adam_eps, + }, + "adamw_config": { + "betas": args.adamw_betas, + "amsgrad": args.adamw_amsgrad, + "eps": args.adamw_eps, + }, + "adamax_config": {"betas": args.adamax_betas, "eps": args.adamax_eps}, + "lbfgs_config": { + "max_iter": args.lbfgs_max_iter, + "max_eval": args.lbfgs_max_eval, + "tolerance_grad": args.lbfgs_tolerance_grad, + "tolerance_change": args.lbfgs_tolerance_change, + "history_size": args.lbfgs_history_size, + "line_search_fn": args.lbfgs_line_search_fn, + }, + "rms_prop_config": { + "alpha": args.rms_prop_alpha, + "eps": args.rms_prop_eps, + "momentum": args.rms_prop_momentum, + "centered": args.rms_prop_centered, + }, + "r_prop_config": { + "etas": args.r_prop_etas, + "step_sizes": args.r_prop_step_sizes, + }, "sgd_config": { "momentum": args.sgd_momentum, "dampening": args.sgd_dampening, "nesterov": args.sgd_nesterov, }, - "adam_config": { - "betas": (0.9, 0.999), - "amsgrad": args.amsgrad, - "eps": args.eps, + "sparse_adam_config": { + "betas": args.sparse_adam_betas, + "eps": args.sparse_adam_eps, + }, + "bert_adam_config": { + "betas": args.bert_adam_betas, + "eps": args.bert_adam_eps, }, - "adamax_config": {"betas": (0.9, 0.999), "eps": args.eps}, - "bert_adam_config": {"betas": (0.9, 0.999), "eps": args.eps}, }, "lr_scheduler_config": { "lr_scheduler": args.lr_scheduler, diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/metrics/test_metrics.py b/tests/metrics/test_metrics.py index 5004324..df2a633 100644 --- a/tests/metrics/test_metrics.py +++ b/tests/metrics/test_metrics.py @@ -1,5 +1,4 @@ import logging -import math import numpy as np @@ -16,26 +15,7 @@ from emmental.metrics.recall import recall_scorer from emmental.metrics.roc_auc import roc_auc_scorer from emmental.metrics.spearman_correlation import spearman_correlation_scorer - - -def isequal(dict_a, dict_b, precision=1e-10): - for key in dict_a: - if ( - key not in dict_b - or abs(dict_a[key] - dict_b[key]) > precision - or (math.isnan(dict_a[key]) and not math.isnan(dict_b[key])) - ): - return False - - for key in dict_b: - if ( - key not in dict_a - or abs(dict_a[key] - dict_b[key]) > precision - or (math.isnan(dict_b[key]) and not math.isnan(dict_a[key])) - ): - return False - - return True +from tests.utils import isequal def test_accuracy(caplog): diff --git a/tests/modules/embeddings.vec b/tests/modules/embeddings.vec new file mode 100644 index 0000000..9303769 --- /dev/null +++ b/tests/modules/embeddings.vec @@ -0,0 +1,4 @@ +3 5 +1 2 2 2 2 2 +3 1 1 1 1 1 +5 3 3 3 3 3 \ No newline at end of file diff --git a/tests/modules/test_embedding_module.py b/tests/modules/test_embedding_module.py new file mode 100644 index 0000000..0f519ae --- /dev/null +++ b/tests/modules/test_embedding_module.py @@ -0,0 +1,61 @@ +import logging + +import torch + +from emmental.modules.embedding_module import EmbeddingModule +from emmental.utils.utils import set_random_seed + +logger = logging.getLogger(__name__) + + +def test_embedding_module(caplog): + """Unit test of Embedding Module""" + + caplog.set_level(logging.INFO) + + # Set random seed seed + set_random_seed(1) + + word_counter = {"1": 1, "2": 3, "3": 1} + weight_tensor = torch.FloatTensor( + [ + [-0.4277, 0.7110, -0.3268, -0.7473, 0.3847], + [0.0000, 0.0000, 0.0000, 0.0000, 0.0000], + [-0.2247, -0.7969, -0.4558, -0.3063, 0.4276], + [2.0000, 2.0000, 2.0000, 2.0000, 2.0000], + [1.0000, 1.0000, 1.0000, 1.0000, 1.0000], + ] + ) + + emb_layer = EmbeddingModule(word_counter=word_counter, word_dim=10, max_size=10) + + assert emb_layer.dim == 10 + # and are default tokens + assert emb_layer.embeddings.weight.size() == (5, 10) + + emb_layer = EmbeddingModule( + word_counter=word_counter, + word_dim=10, + embedding_file="tests/modules/embeddings.vec", + fix_emb=True, + ) + + assert emb_layer.dim == 5 + assert emb_layer.embeddings.weight.size() == (5, 5) + assert torch.max(torch.abs(emb_layer.embeddings.weight.data - weight_tensor)) < 1e-4 + + assert ( + torch.max( + torch.abs(emb_layer(torch.LongTensor([1, 2])) - weight_tensor[1:3, :]) + ) + < 1e-4 + ) + + # With threshold + word_counter = {"1": 3, "2": 1, "3": 1} + emb_layer = EmbeddingModule(word_counter=word_counter, word_dim=10, threshold=2) + assert emb_layer.embeddings.weight.size() == (3, 10) + + # No word counter + emb_layer = EmbeddingModule(embedding_file="tests/modules/embeddings.vec") + assert emb_layer.embeddings.weight.size() == (5, 5) diff --git a/tests/modules/test_identity_module.py b/tests/modules/test_identity_module.py new file mode 100644 index 0000000..55109eb --- /dev/null +++ b/tests/modules/test_identity_module.py @@ -0,0 +1,18 @@ +import logging + +import torch + +from emmental.modules.identity_module import IdentityModule + +logger = logging.getLogger(__name__) + + +def test_identity_module(caplog): + """Unit test of Identity Module""" + + caplog.set_level(logging.INFO) + + identity_module = IdentityModule() + + input = torch.randn(10, 10) + assert torch.equal(input, identity_module(input)) diff --git a/tests/modules/test_rnn_module.py b/tests/modules/test_rnn_module.py new file mode 100644 index 0000000..f37dd6e --- /dev/null +++ b/tests/modules/test_rnn_module.py @@ -0,0 +1,58 @@ +import logging + +import torch + +from emmental.modules.rnn_module import RNN +from emmental.utils.utils import pad_batch + +logger = logging.getLogger(__name__) + + +def test_rnn_module(caplog): + """Unit test of RNN Module""" + + caplog.set_level(logging.INFO) + + n_class = 2 + emb_size = 10 + lstm_hidden = 20 + batch_size = 3 + seq_len = 4 + + # Single direction RNN + rnn = RNN( + num_classes=n_class, + emb_size=emb_size, + lstm_hidden=lstm_hidden, + attention=True, + dropout=0.2, + bidirectional=False, + ) + _, input_mask = pad_batch(torch.randn(batch_size, seq_len)) + + assert rnn(torch.randn(batch_size, seq_len, emb_size)).size() == (3, n_class) + assert rnn(torch.randn(batch_size, seq_len, emb_size), input_mask).size() == ( + 3, + n_class, + ) + + # Bi-direction RNN + rnn = RNN( + num_classes=0, + emb_size=emb_size, + lstm_hidden=lstm_hidden, + attention=False, + dropout=0.2, + bidirectional=True, + ) + + _, input_mask = pad_batch(torch.randn(batch_size, seq_len)) + + assert rnn(torch.randn(batch_size, seq_len, emb_size)).size() == ( + 3, + 2 * lstm_hidden, + ) + assert rnn(torch.randn(batch_size, seq_len, emb_size), input_mask).size() == ( + 3, + 2 * lstm_hidden, + ) diff --git a/tests/modules/test_sparse_linear_module.py b/tests/modules/test_sparse_linear_module.py new file mode 100644 index 0000000..3754bae --- /dev/null +++ b/tests/modules/test_sparse_linear_module.py @@ -0,0 +1,29 @@ +import logging + +import numpy as np +import torch + +from emmental.modules.sparse_linear_module import SparseLinear + +logger = logging.getLogger(__name__) + + +def test_sparse_linear_module(caplog): + """Unit test of Sparse Linear Module""" + + caplog.set_level(logging.INFO) + + # Test w/o bias + sparse_linear_module = SparseLinear(10, 2) + + index = torch.from_numpy(np.random.randint(10, size=100)).view(10, 10) + weight = torch.randn(10, 10) + + assert isinstance(sparse_linear_module(index, weight), torch.Tensor) + assert sparse_linear_module(index, weight).size() == (10, 2) + + # Test with bias + sparse_linear_module = SparseLinear(10, 2, bias=True) + + assert isinstance(sparse_linear_module(index, weight), torch.Tensor) + assert sparse_linear_module(index, weight).size() == (10, 2) diff --git a/tests/optimizers/__init__.py b/tests/optimizers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/optimizers/test_adadelta_optimizer.py b/tests/optimizers/test_adadelta_optimizer.py new file mode 100644 index 0000000..fd1f9db --- /dev/null +++ b/tests/optimizers/test_adadelta_optimizer.py @@ -0,0 +1,56 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner +from tests.utils import isequal + +logger = logging.getLogger(__name__) + + +def test_adadelta_optimizer(caplog): + """Unit test of Adadelta optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "adadelta" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default Adadelta setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert isequal( + emmental_learner.optimizer.defaults, + {"lr": 0.001, "rho": 0.9, "eps": 1e-06, "weight_decay": 0}, + ) + + # Test new Adadelta setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": {"rho": 0.6, "eps": 1e-05}, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert isequal( + emmental_learner.optimizer.defaults, + {"lr": 0.02, "rho": 0.6, "eps": 1e-05, "weight_decay": 0.05}, + ) + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_adagrad_optimizer.py b/tests/optimizers/test_adagrad_optimizer.py new file mode 100644 index 0000000..7113675 --- /dev/null +++ b/tests/optimizers/test_adagrad_optimizer.py @@ -0,0 +1,72 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner +from tests.utils import isequal + +logger = logging.getLogger(__name__) + + +def test_adagrad_optimizer(caplog): + """Unit test of Adagrad optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "adagrad" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default Adagrad setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert isequal( + emmental_learner.optimizer.defaults, + { + "lr": 0.001, + "lr_decay": 0, + "initial_accumulator_value": 0, + "eps": 1e-10, + "weight_decay": 0, + }, + ) + + # Test new Adagrad setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": { + "lr_decay": 0.1, + "initial_accumulator_value": 0.2, + "eps": 1e-5, + }, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert isequal( + emmental_learner.optimizer.defaults, + { + "lr": 0.02, + "lr_decay": 0.1, + "initial_accumulator_value": 0.2, + "eps": 1e-5, + "weight_decay": 0.05, + }, + ) + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_adam_optimizer.py b/tests/optimizers/test_adam_optimizer.py new file mode 100644 index 0000000..b60cc09 --- /dev/null +++ b/tests/optimizers/test_adam_optimizer.py @@ -0,0 +1,65 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_adam_optimizer(caplog): + """Unit test of Adam optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "adam" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default Adam setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "betas": (0.9, 0.999), + "eps": 1e-08, + "amsgrad": False, + "weight_decay": 0, + } + + # Test new Adam setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": { + "betas": (0.9, 0.99), + "eps": 1e-05, + "amsgrad": True, + }, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "betas": (0.9, 0.99), + "eps": 1e-05, + "amsgrad": True, + "weight_decay": 0.05, + } + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_adamax_optimizer.py b/tests/optimizers/test_adamax_optimizer.py new file mode 100644 index 0000000..0a6123d --- /dev/null +++ b/tests/optimizers/test_adamax_optimizer.py @@ -0,0 +1,59 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_adamax_optimizer(caplog): + """Unit test of Adamax optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "adamax" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default Adamax setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "betas": (0.9, 0.999), + "eps": 1e-08, + "weight_decay": 0, + } + + # Test new Adamax setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": {"betas": (0.9, 0.99), "eps": 1e-05}, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "betas": (0.9, 0.99), + "eps": 1e-05, + "weight_decay": 0.05, + } + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_adamw_optimizer.py b/tests/optimizers/test_adamw_optimizer.py new file mode 100644 index 0000000..e77c3d5 --- /dev/null +++ b/tests/optimizers/test_adamw_optimizer.py @@ -0,0 +1,65 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_adamw_optimizer(caplog): + """Unit test of AdamW optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "adamw" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default AdamW setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "betas": (0.9, 0.999), + "eps": 1e-08, + "amsgrad": False, + "weight_decay": 0, + } + + # Test new AdamW setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": { + "betas": (0.9, 0.99), + "eps": 1e-05, + "amsgrad": True, + }, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "betas": (0.9, 0.99), + "eps": 1e-05, + "amsgrad": True, + "weight_decay": 0.05, + } + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_asgd_optimizer.py b/tests/optimizers/test_asgd_optimizer.py new file mode 100644 index 0000000..62260ef --- /dev/null +++ b/tests/optimizers/test_asgd_optimizer.py @@ -0,0 +1,62 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner +from tests.utils import isequal + +logger = logging.getLogger(__name__) + + +def test_asgd_optimizer(caplog): + """Unit test of ASGD optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "asgd" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default ASGD setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert isequal( + emmental_learner.optimizer.defaults, + { + "lr": 0.001, + "lambd": 0.0001, + "alpha": 0.75, + "t0": 1_000_000.0, + "weight_decay": 0, + }, + ) + + # Test default ASGD setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": {"lambd": 0.1, "alpha": 0.5, "t0": 1e5}, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert isequal( + emmental_learner.optimizer.defaults, + {"lr": 0.02, "lambd": 0.1, "alpha": 0.5, "t0": 1e5, "weight_decay": 0.05}, + ) + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_bert_adam_optimizer.py b/tests/optimizers/test_bert_adam_optimizer.py new file mode 100644 index 0000000..133f6ea --- /dev/null +++ b/tests/optimizers/test_bert_adam_optimizer.py @@ -0,0 +1,67 @@ +import logging +import shutil + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_bert_adam_optimizer(caplog): + """Unit test of BertAdam optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "bert_adam" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default BertAdam setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "betas": (0.9, 0.999), + "eps": 1e-08, + "weight_decay": 0.0, + } + + # Test new BertAdam setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": {"betas": (0.8, 0.9), "eps": 1e-05}, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "betas": (0.8, 0.9), + "eps": 1e-05, + "weight_decay": 0.05, + } + + # Test BertAdam setp + emmental_learner.optimizer.zero_grad() + torch.Tensor(1) + F.mse_loss(model(torch.randn(1, 1)), torch.randn(1, 1)).backward() + emmental_learner.optimizer.step() + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_lbfgs_optimizer.py b/tests/optimizers/test_lbfgs_optimizer.py new file mode 100644 index 0000000..d46ff53 --- /dev/null +++ b/tests/optimizers/test_lbfgs_optimizer.py @@ -0,0 +1,72 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_lbfgs_optimizer(caplog): + """Unit test of LBFGS optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "lbfgs" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default LBFGS setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "max_iter": 20, + "max_eval": 25, + "tolerance_grad": 1e-07, + "tolerance_change": 1e-09, + "history_size": 100, + "line_search_fn": None, + } + + # Test new LBFGS setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": { + "max_iter": 30, + "max_eval": 40, + "tolerance_grad": 1e-04, + "tolerance_change": 1e-05, + "history_size": 10, + "line_search_fn": "strong_wolfe", + }, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "max_iter": 30, + "max_eval": 40, + "tolerance_grad": 1e-04, + "tolerance_change": 1e-05, + "history_size": 10, + "line_search_fn": "strong_wolfe", + } + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_r_prop_optimizer.py b/tests/optimizers/test_r_prop_optimizer.py new file mode 100644 index 0000000..0f6090e --- /dev/null +++ b/tests/optimizers/test_r_prop_optimizer.py @@ -0,0 +1,57 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_r_prop_optimizer(caplog): + """Unit test of Rprop optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "r_prop" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default Rprop setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "etas": (0.5, 1.2), + "step_sizes": (1e-06, 50), + } + + # Test new Rprop setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": {"etas": (0.3, 1.5), "step_sizes": (1e-04, 30)}, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "etas": (0.3, 1.5), + "step_sizes": (1e-04, 30), + } + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_rms_prop_optimizer.py b/tests/optimizers/test_rms_prop_optimizer.py new file mode 100644 index 0000000..928dea0 --- /dev/null +++ b/tests/optimizers/test_rms_prop_optimizer.py @@ -0,0 +1,68 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_rms_prop_optimizer(caplog): + """Unit test of RMSprop optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "rms_prop" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default RMSprop setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "alpha": 0.99, + "eps": 1e-08, + "momentum": 0, + "centered": False, + "weight_decay": 0, + } + + # Test new RMSprop setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": { + "alpha": 0.9, + "eps": 1e-05, + "momentum": 0.1, + "centered": True, + }, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "alpha": 0.9, + "eps": 1e-05, + "momentum": 0.1, + "centered": True, + "weight_decay": 0.05, + } + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_sgd_optimizer.py b/tests/optimizers/test_sgd_optimizer.py new file mode 100644 index 0000000..90dd321 --- /dev/null +++ b/tests/optimizers/test_sgd_optimizer.py @@ -0,0 +1,65 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_sgd_optimizer(caplog): + """Unit test of SGD optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "sgd" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default SGD setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "momentum": 0, + "dampening": 0, + "nesterov": False, + "weight_decay": 0.0, + } + + # Test new SGD setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": { + "momentum": 0.1, + "dampening": 0, + "nesterov": True, + }, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "momentum": 0.1, + "dampening": 0, + "nesterov": True, + "weight_decay": 0.05, + } + + shutil.rmtree(dirpath) diff --git a/tests/optimizers/test_sparse_adam_optimizer.py b/tests/optimizers/test_sparse_adam_optimizer.py new file mode 100644 index 0000000..fb79aa5 --- /dev/null +++ b/tests/optimizers/test_sparse_adam_optimizer.py @@ -0,0 +1,57 @@ +import logging +import shutil + +import torch.nn as nn + +import emmental +from emmental import Meta +from emmental.learner import EmmentalLearner + +logger = logging.getLogger(__name__) + + +def test_sparse_adam_optimizer(caplog): + """Unit test of SparseAdam optimizer""" + + caplog.set_level(logging.INFO) + + optimizer = "sparse_adam" + dirpath = "temp_test_optimizer" + model = nn.Linear(1, 1) + emmental_learner = EmmentalLearner() + + Meta.reset() + emmental.init(dirpath) + + # Test default SparseAdam setting + config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.001, + "betas": (0.9, 0.999), + "eps": 1e-08, + } + + # Test new SparseAdam setting + config = { + "learner_config": { + "optimizer_config": { + "optimizer": optimizer, + "lr": 0.02, + "l2": 0.05, + f"{optimizer}_config": {"betas": (0.8, 0.9), "eps": 1e-05}, + } + } + } + emmental.Meta.update_config(config) + emmental_learner._set_optimizer(model) + + assert emmental_learner.optimizer.defaults == { + "lr": 0.02, + "betas": (0.8, 0.9), + "eps": 1e-05, + } + + shutil.rmtree(dirpath) diff --git a/tests/schedulers/test_mixed_scheduler.py b/tests/schedulers/test_mixed_scheduler.py new file mode 100644 index 0000000..8dd624f --- /dev/null +++ b/tests/schedulers/test_mixed_scheduler.py @@ -0,0 +1,67 @@ +import logging + +import numpy as np +import torch + +import emmental +from emmental.data import EmmentalDataLoader, EmmentalDataset +from emmental.schedulers.mixed_scheduler import MixedScheduler + +logger = logging.getLogger(__name__) + + +def test_mixed_scheduler(caplog): + """Unit test of mixed scheduler""" + + caplog.set_level(logging.INFO) + + emmental.Meta.init() + + task1 = "task1" + x1 = np.random.rand(20, 2) + y1 = torch.from_numpy(np.random.rand(20)) + + task2 = "task2" + x2 = np.random.rand(30, 3) + y2 = torch.from_numpy(np.random.rand(30)) + + dataloaders = [ + EmmentalDataLoader( + task_to_label_dict={task_name: "label"}, + dataset=EmmentalDataset( + name=task_name, X_dict={"feature": x}, Y_dict={"label": y} + ), + split="train", + batch_size=10, + shuffle=True, + ) + for task_name, x, y in [(task1, x1, y1), (task2, x2, y2)] + ] + + scheduler = MixedScheduler() + + assert scheduler.get_num_batches(dataloaders) == 2 + + batch_task_names_1 = [ + batch_data[0][-2] for batch_data in scheduler.get_batches(dataloaders) + ] + batch_task_names_2 = [ + batch_data[1][-2] for batch_data in scheduler.get_batches(dataloaders) + ] + + assert batch_task_names_1 == [task1, task1] + assert batch_task_names_2 == [task2, task2] + + scheduler = MixedScheduler(fillup=True) + + assert scheduler.get_num_batches(dataloaders) == 3 + + batch_task_names_1 = [ + batch_data[0][-2] for batch_data in scheduler.get_batches(dataloaders) + ] + batch_task_names_2 = [ + batch_data[1][-2] for batch_data in scheduler.get_batches(dataloaders) + ] + + assert batch_task_names_1 == [task1, task1, task1] + assert batch_task_names_2 == [task2, task2, task2] diff --git a/tests/schedulers/test_round_robin_scheduler.py b/tests/schedulers/test_round_robin_scheduler.py new file mode 100644 index 0000000..250aada --- /dev/null +++ b/tests/schedulers/test_round_robin_scheduler.py @@ -0,0 +1,63 @@ +import logging + +import numpy as np +import torch + +import emmental +from emmental.data import EmmentalDataLoader, EmmentalDataset +from emmental.schedulers.round_robin_scheduler import RoundRobinScheduler +from emmental.utils.utils import set_random_seed + +logger = logging.getLogger(__name__) + + +def test_round_robin_scheduler(caplog): + """Unit test of round robin scheduler""" + + caplog.set_level(logging.INFO) + + emmental.Meta.init() + + # Set random seed seed + set_random_seed(2) + + task1 = "task1" + x1 = np.random.rand(20, 2) + y1 = torch.from_numpy(np.random.rand(20)) + + task2 = "task2" + x2 = np.random.rand(30, 3) + y2 = torch.from_numpy(np.random.rand(30)) + + dataloaders = [ + EmmentalDataLoader( + task_to_label_dict={task_name: "label"}, + dataset=EmmentalDataset( + name=task_name, X_dict={"feature": x}, Y_dict={"label": y} + ), + split="train", + batch_size=10, + shuffle=True, + ) + for task_name, x, y in [(task1, x1, y1), (task2, x2, y2)] + ] + + scheduler = RoundRobinScheduler() + + assert scheduler.get_num_batches(dataloaders) == 5 + + batch_task_names = [ + batch_data[-2] for batch_data in scheduler.get_batches(dataloaders) + ] + + assert batch_task_names == [task2, task1, task2, task2, task1] + + scheduler = RoundRobinScheduler(fillup=True) + + assert scheduler.get_num_batches(dataloaders) == 6 + + batch_task_names = [ + batch_data[-2] for batch_data in scheduler.get_batches(dataloaders) + ] + + assert batch_task_names == [task2, task1, task2, task2, task1, task1] diff --git a/tests/schedulers/test_sequential_scheduler.py b/tests/schedulers/test_sequential_scheduler.py new file mode 100644 index 0000000..f613180 --- /dev/null +++ b/tests/schedulers/test_sequential_scheduler.py @@ -0,0 +1,59 @@ +import logging + +import numpy as np +import torch + +import emmental +from emmental.data import EmmentalDataLoader, EmmentalDataset +from emmental.schedulers.sequential_scheduler import SequentialScheduler + +logger = logging.getLogger(__name__) + + +def test_sequential_scheduler(caplog): + """Unit test of sequential scheduler""" + + caplog.set_level(logging.INFO) + + emmental.Meta.init() + + task1 = "task1" + x1 = np.random.rand(20, 2) + y1 = torch.from_numpy(np.random.rand(20)) + + task2 = "task2" + x2 = np.random.rand(30, 3) + y2 = torch.from_numpy(np.random.rand(30)) + + dataloaders = [ + EmmentalDataLoader( + task_to_label_dict={task_name: "label"}, + dataset=EmmentalDataset( + name=task_name, X_dict={"feature": x}, Y_dict={"label": y} + ), + split="train", + batch_size=10, + shuffle=True, + ) + for task_name, x, y in [(task1, x1, y1), (task2, x2, y2)] + ] + + scheduler = SequentialScheduler() + + assert scheduler.get_num_batches(dataloaders) == 5 + + batch_task_names = [ + batch_data[-2] for batch_data in scheduler.get_batches(dataloaders) + ] + + assert batch_task_names == [task1, task1, task2, task2, task2] + + scheduler = SequentialScheduler(fillup=True) + + assert scheduler.get_num_batches(dataloaders) == 6 + + batch_task_names = [ + batch_data[-2] for batch_data in scheduler.get_batches(dataloaders) + ] + + assert batch_task_names == [task1, task1, task1, task2, task2, task2] diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..e5fcfa7 --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,100 @@ +import logging +import shutil +from functools import partial + +import torch.nn as nn +import torch.nn.functional as F + +import emmental +from emmental import Meta +from emmental.model import EmmentalModel +from emmental.scorer import Scorer +from emmental.task import EmmentalTask + +logger = logging.getLogger(__name__) + + +def test_model(caplog): + """Unit test of model""" + + caplog.set_level(logging.INFO) + + dirpath = "temp_test_model" + + Meta.reset() + emmental.init(dirpath) + + def ce_loss(module_name, immediate_output_dict, Y, active): + return F.cross_entropy( + immediate_output_dict[module_name][0][active], (Y.view(-1))[active] + ) + + def output(module_name, immediate_output_dict): + return F.softmax(immediate_output_dict[module_name][0], dim=1) + + task1 = EmmentalTask( + name="task_1", + module_pool=nn.ModuleDict( + {"m1": nn.Linear(10, 10, bias=False), "m2": nn.Linear(10, 2, bias=False)} + ), + task_flow=[ + {"name": "m1", "module": "m1", "inputs": [("_input_", "data")]}, + {"name": "m2", "module": "m2", "inputs": [("m1", 0)]}, + ], + loss_func=partial(ce_loss, "m2"), + output_func=partial(output, "m2"), + scorer=Scorer(metrics=["accuracy"]), + ) + + model = EmmentalModel(name="test", tasks=task1) + + assert model.name == "test" + assert model.task_names == set(["task_1"]) + assert model.module_pool["m1"].module.weight.data.size() == (10, 10) + assert model.module_pool["m2"].module.weight.data.size() == (2, 10) + + task1 = EmmentalTask( + name="task_1", + module_pool=nn.ModuleDict( + {"m1": nn.Linear(10, 5, bias=False), "m2": nn.Linear(5, 2, bias=False)} + ), + task_flow=[ + {"name": "m1", "module": "m1", "inputs": [("_input_", "data")]}, + {"name": "m2", "module": "m2", "inputs": [("m1", 0)]}, + ], + loss_func=partial(ce_loss, "m2"), + output_func=partial(output, "m2"), + scorer=Scorer(metrics=["accuracy"]), + ) + + model.update_task(task1) + + assert model.module_pool["m1"].module.weight.data.size() == (5, 10) + assert model.module_pool["m2"].module.weight.data.size() == (2, 5) + + task2 = EmmentalTask( + name="task_2", + module_pool=nn.ModuleDict( + {"m1": nn.Linear(10, 5, bias=False), "m2": nn.Linear(5, 2, bias=False)} + ), + task_flow=[ + {"name": "m1", "module": "m1", "inputs": [("_input_", "data")]}, + {"name": "m2", "module": "m2", "inputs": [("m1", 0)]}, + ], + loss_func=partial(ce_loss, "m2"), + output_func=partial(output, "m2"), + scorer=Scorer(metrics=["accuracy"]), + ) + + model.add_task(task2) + + assert model.task_names == set(["task_1", "task_2"]) + + model.remove_task("task_1") + assert model.task_names == set(["task_2"]) + + model.save(f"{dirpath}/saved_model.pth") + + model.load(f"{dirpath}/saved_model.pth") + + shutil.rmtree(dirpath) diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 0000000..cf69c02 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,21 @@ +import math + + +def isequal(dict_a, dict_b, precision=1e-10): + for key in dict_a: + if ( + key not in dict_b + or abs(dict_a[key] - dict_b[key]) > precision + or (math.isnan(dict_a[key]) and not math.isnan(dict_b[key])) + ): + return False + + for key in dict_b: + if ( + key not in dict_a + or abs(dict_a[key] - dict_b[key]) > precision + or (math.isnan(dict_b[key]) and not math.isnan(dict_a[key])) + ): + return False + + return True diff --git a/tests/utils/test_parse_args.py b/tests/utils/test_parse_args.py new file mode 100644 index 0000000..b417a9c --- /dev/null +++ b/tests/utils/test_parse_args.py @@ -0,0 +1,103 @@ +import logging + +from emmental.utils.parse_args import parse_args, parse_args_to_config + +logger = logging.getLogger(__name__) + + +def test_parse_args(caplog): + """Unit test of parsing args""" + + caplog.set_level(logging.INFO) + + parser = parse_args() + args = parser.parse_args(["--seed", "0"]) + assert args.seed == 0 + + config = parse_args_to_config(args) + + assert config == { + "meta_config": {"seed": 0, "verbose": True, "log_path": "logs"}, + "data_config": {"min_data_len": 0, "max_data_len": 0}, + "model_config": {"model_path": None, "device": 0, "dataparallel": True}, + "learner_config": { + "fp16": False, + "n_epochs": 3, + "train_split": ["train"], + "valid_split": ["valid"], + "test_split": ["test"], + "ignore_index": None, + "optimizer_config": { + "optimizer": "adam", + "lr": 0.001, + "l2": 0.0, + "grad_clip": None, + "asgd_config": {"lambd": 0.0001, "alpha": 0.75, "t0": 1000000.0}, + "adadelta_config": {"rho": 0.9, "eps": 1e-06}, + "adagrad_config": { + "lr_decay": 0, + "initial_accumulator_value": 0, + "eps": 1e-10, + }, + "adam_config": {"betas": (0.9, 0.999), "amsgrad": False, "eps": 1e-08}, + "adamw_config": {"betas": (0.9, 0.999), "amsgrad": False, "eps": 1e-08}, + "adamax_config": {"betas": (0.9, 0.999), "eps": 1e-08}, + "lbfgs_config": { + "max_iter": 20, + "max_eval": None, + "tolerance_grad": 1e-07, + "tolerance_change": 1e-09, + "history_size": 100, + "line_search_fn": None, + }, + "rms_prop_config": { + "alpha": 0.99, + "eps": 1e-08, + "momentum": 0, + "centered": False, + }, + "r_prop_config": {"etas": (0.5, 1.2), "step_sizes": (1e-06, 50)}, + "sgd_config": {"momentum": 0.0, "dampening": 0, "nesterov": False}, + "sparse_adam_config": {"betas": (0.9, 0.999), "eps": 1e-08}, + "bert_adam_config": {"betas": (0.9, 0.999), "eps": 1e-08}, + }, + "lr_scheduler_config": { + "lr_scheduler": None, + "warmup_steps": None, + "warmup_unit": "batch", + "warmup_percentage": None, + "min_lr": 0.0, + "linear_config": {"min_lr": 0.0}, + "exponential_config": {"gamma": 0.9}, + "plateau_config": {"factor": 0.5, "patience": 10, "threshold": 0.0001}, + "step_config": {"step_size": 1, "gamma": 0.01, "last_epoch": -1}, + "multi_step_config": { + "milestones": [10000], + "gamma": 0.01, + "last_epoch": -1, + }, + "cosine_annealing_config": {"last_epoch": -1}, + }, + "task_scheduler_config": { + "task_scheduler": "round_robin", + "sequential_scheduler_config": {"fillup": False}, + "round_robin_scheduler_config": {"fillup": False}, + "mixed_scheduler_config": {"fillup": False}, + }, + }, + "logging_config": { + "counter_unit": "epoch", + "evaluation_freq": 1, + "writer_config": {"writer": "tensorboard", "verbose": True}, + "checkpointing": False, + "checkpointer_config": { + "checkpoint_path": None, + "checkpoint_freq": 1, + "checkpoint_metric": {"model/train/all/loss": "min"}, + "checkpoint_task_metrics": None, + "checkpoint_runway": 0, + "clear_intermediate_checkpoints": True, + "clear_all_checkpoints": False, + }, + }, + }