mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,22 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
from nemo.utils.app_state import AppState
|
||||
from nemo.utils.nemo_logging import Logger as _Logger
|
||||
from nemo.utils.nemo_logging import LogMode as logging_mode
|
||||
from nemo.utils.lightning_logger_patch import add_memory_handlers_to_pl_logger
|
||||
|
||||
logging = _Logger()
|
||||
add_memory_handlers_to_pl_logger()
|
||||
@@ -0,0 +1,386 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from threading import Lock
|
||||
from typing import Dict, Optional
|
||||
|
||||
from nemo.utils.metaclasses import Singleton
|
||||
|
||||
|
||||
@dataclass()
|
||||
class ModelMetadataRegistry:
|
||||
guid: str
|
||||
gidx: int
|
||||
restoration_path: Optional[str] = None
|
||||
|
||||
|
||||
class AppState(metaclass=Singleton):
|
||||
def __init__(self):
|
||||
# method call lock
|
||||
self.__lock = Lock()
|
||||
|
||||
# TODO: should we store global config in hydra_runner?
|
||||
self._app_cfg = None
|
||||
|
||||
# World info
|
||||
self._device_id = None
|
||||
self._local_rank = None
|
||||
self._global_rank = None
|
||||
self._model_parallel_rank = None
|
||||
self._data_parallel_rank = None
|
||||
|
||||
self._world_size = None
|
||||
self._model_parallel_size = None
|
||||
self._model_parallel_group = None
|
||||
self._data_parallel_size = None
|
||||
self._data_parallel_group = None
|
||||
|
||||
self._random_seed = None
|
||||
|
||||
# Logging info
|
||||
self._log_dir = None
|
||||
self._exp_dir = None
|
||||
self._name = None
|
||||
self._checkpoint_name = None
|
||||
self._version = None
|
||||
self._create_checkpoint_callback = None
|
||||
self._checkpoint_callback_params = None
|
||||
|
||||
# Save and Restore (.nemo)
|
||||
self._tmpdir_name = None
|
||||
self._model_config_yaml = "model_config.yaml"
|
||||
self._model_weights_ckpt = "model_weights.ckpt"
|
||||
self._model_restore_path = None
|
||||
self._all_model_restore_paths = []
|
||||
self._model_guid_map = {} # type: Dict[str, ModelMetadataRegistry]
|
||||
|
||||
@property
|
||||
def device_id(self):
|
||||
""" Property returns the device_id
|
||||
Returns:
|
||||
device_id
|
||||
"""
|
||||
return self._device_id
|
||||
|
||||
@device_id.setter
|
||||
def device_id(self, id):
|
||||
""" Property sets the device_id.
|
||||
Args:
|
||||
size (int): The device id.
|
||||
"""
|
||||
self._device_id = id
|
||||
|
||||
@property
|
||||
def world_size(self):
|
||||
""" Property returns the total number of GPUs.
|
||||
Returns:
|
||||
Total number of GPUs.
|
||||
"""
|
||||
return self._world_size
|
||||
|
||||
@world_size.setter
|
||||
def world_size(self, size):
|
||||
""" Property sets the total number of GPUs.
|
||||
Args:
|
||||
size (int): Total number of GPUs.
|
||||
"""
|
||||
self._world_size = size
|
||||
|
||||
@property
|
||||
def model_parallel_size(self):
|
||||
""" Property returns the number of GPUs in each model parallel group.
|
||||
Returns:
|
||||
Number of GPUs in each model parallel group.
|
||||
"""
|
||||
return self._model_parallel_size
|
||||
|
||||
@model_parallel_size.setter
|
||||
def model_parallel_size(self, size):
|
||||
""" Property sets the number of GPUs in each model parallel group.
|
||||
Args:
|
||||
size (int): Number of GPUs in each model parallel group.
|
||||
"""
|
||||
self._model_parallel_size = size
|
||||
|
||||
@property
|
||||
def data_parallel_size(self):
|
||||
""" Property returns the number of GPUs in each data parallel group.
|
||||
Returns:
|
||||
Number of GPUs in each data parallel group.
|
||||
"""
|
||||
return self._data_parallel_size
|
||||
|
||||
@data_parallel_size.setter
|
||||
def data_parallel_size(self, size):
|
||||
""" Property sets the number of GPUs in each data parallel group.
|
||||
Args:
|
||||
size (int): Number of GPUs in each data parallel group.
|
||||
"""
|
||||
self._data_parallel_size = size
|
||||
|
||||
@property
|
||||
def local_rank(self):
|
||||
""" Property returns the local rank.
|
||||
Returns:
|
||||
Local rank.
|
||||
"""
|
||||
return self._local_rank
|
||||
|
||||
@local_rank.setter
|
||||
def local_rank(self, rank):
|
||||
""" Property sets the local rank.
|
||||
Args:
|
||||
rank (int): Local rank.
|
||||
"""
|
||||
self._local_rank = rank
|
||||
|
||||
@property
|
||||
def global_rank(self):
|
||||
""" Property returns the global rank.
|
||||
Returns:
|
||||
Global rank.
|
||||
"""
|
||||
return self._global_rank
|
||||
|
||||
@global_rank.setter
|
||||
def global_rank(self, rank):
|
||||
""" Property sets the global rank.
|
||||
Args:
|
||||
rank (int): Global rank.
|
||||
"""
|
||||
self._global_rank = rank
|
||||
|
||||
@property
|
||||
def model_parallel_rank(self):
|
||||
""" Property returns the model parallel rank.
|
||||
Returns:
|
||||
Model parallel rank.
|
||||
"""
|
||||
return self._model_parallel_rank
|
||||
|
||||
@model_parallel_rank.setter
|
||||
def model_parallel_rank(self, rank):
|
||||
""" Property sets the model parallel rank.
|
||||
Args:
|
||||
rank (int): Model parallel rank.
|
||||
"""
|
||||
self._model_parallel_rank = rank
|
||||
|
||||
@property
|
||||
def model_parallel_group(self):
|
||||
""" Property returns the model parallel group.
|
||||
Returns:
|
||||
Model parallel group.
|
||||
"""
|
||||
return self._model_parallel_group
|
||||
|
||||
@model_parallel_group.setter
|
||||
def model_parallel_group(self, group):
|
||||
""" Property sets the model parallel group.
|
||||
Args:
|
||||
group: Model parallel group.
|
||||
"""
|
||||
self._model_parallel_group = group
|
||||
|
||||
@property
|
||||
def data_parallel_rank(self):
|
||||
""" Property returns the data parallel rank.
|
||||
Returns:
|
||||
Data parallel rank.
|
||||
"""
|
||||
return self._data_parallel_rank
|
||||
|
||||
@data_parallel_rank.setter
|
||||
def data_parallel_rank(self, rank):
|
||||
""" Property sets the data parallel rank.
|
||||
Args:
|
||||
rank (int): Data parallel rank.
|
||||
"""
|
||||
self._data_parallel_rank = rank
|
||||
|
||||
@property
|
||||
def data_parallel_group(self):
|
||||
""" Property returns the data parallel group.
|
||||
Returns:
|
||||
Data parallel group.
|
||||
"""
|
||||
return self._data_parallel_group
|
||||
|
||||
@data_parallel_group.setter
|
||||
def data_parallel_group(self, group):
|
||||
""" Property sets the data parallel group.
|
||||
Args:
|
||||
group: Data parallel group.
|
||||
"""
|
||||
self._data_parallel_group = group
|
||||
|
||||
@property
|
||||
def random_seed(self):
|
||||
""" Property returns the random seed.
|
||||
Returns:
|
||||
Random seed.
|
||||
"""
|
||||
return self._random_seed
|
||||
|
||||
@random_seed.setter
|
||||
def random_seed(self, seed):
|
||||
""" Property sets the random seed.
|
||||
Args:
|
||||
seed (int): Random seed.
|
||||
"""
|
||||
self._random_seed = seed
|
||||
|
||||
@property
|
||||
def log_dir(self):
|
||||
"""Returns the log_dir set by exp_manager.
|
||||
"""
|
||||
return self._log_dir
|
||||
|
||||
@log_dir.setter
|
||||
def log_dir(self, dir):
|
||||
"""Sets the log_dir property.
|
||||
|
||||
Args:
|
||||
dir (str): Log_dir set by exp_manager.
|
||||
"""
|
||||
self._log_dir = dir
|
||||
|
||||
@property
|
||||
def exp_dir(self):
|
||||
"""Returns the exp_dir set by exp_manager.
|
||||
"""
|
||||
return self._exp_dir
|
||||
|
||||
@exp_dir.setter
|
||||
def exp_dir(self, dir):
|
||||
"""Sets the log_dir property.
|
||||
|
||||
Args:
|
||||
dir (str): Log_dir set by exp_manager.
|
||||
"""
|
||||
self._exp_dir = dir
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
"""Returns the name set by exp_manager.
|
||||
"""
|
||||
return self._name
|
||||
|
||||
@name.setter
|
||||
def name(self, name):
|
||||
"""Sets the name property.
|
||||
|
||||
Args:
|
||||
dir (str): name set by exp_manager.
|
||||
"""
|
||||
self._name = name
|
||||
|
||||
@property
|
||||
def checkpoint_name(self):
|
||||
"""Returns the name set by exp_manager.
|
||||
"""
|
||||
return self._checkpoint_name
|
||||
|
||||
@checkpoint_name.setter
|
||||
def checkpoint_name(self, name):
|
||||
"""Sets the name property.
|
||||
|
||||
Args:
|
||||
dir (str): name set by exp_manager.
|
||||
"""
|
||||
self._checkpoint_name = name
|
||||
|
||||
@property
|
||||
def version(self):
|
||||
"""Returns the version set by exp_manager.
|
||||
"""
|
||||
return self._version
|
||||
|
||||
@version.setter
|
||||
def version(self, version):
|
||||
"""Sets the version property.
|
||||
|
||||
Args:
|
||||
dir (str): version set by exp_manager.
|
||||
"""
|
||||
self._version = version
|
||||
|
||||
@property
|
||||
def create_checkpoint_callback(self):
|
||||
"""Returns the create_checkpoint_callback set by exp_manager.
|
||||
"""
|
||||
return self._create_checkpoint_callback
|
||||
|
||||
@create_checkpoint_callback.setter
|
||||
def create_checkpoint_callback(self, create_checkpoint_callback):
|
||||
"""Sets the create_checkpoint_callback property.
|
||||
|
||||
Args:
|
||||
dir (bool): create_checkpoint_callback set by exp_manager.
|
||||
"""
|
||||
self._create_checkpoint_callback = create_checkpoint_callback
|
||||
|
||||
@property
|
||||
def checkpoint_callback_params(self):
|
||||
"""Returns the version set by exp_manager.
|
||||
"""
|
||||
return self._checkpoint_callback_params
|
||||
|
||||
@checkpoint_callback_params.setter
|
||||
def checkpoint_callback_params(self, params):
|
||||
"""Sets the name property.
|
||||
|
||||
Args:
|
||||
params (dict): checkpoint_callback_params set by exp_manager.
|
||||
"""
|
||||
self._checkpoint_callback_params = params
|
||||
|
||||
@property
|
||||
def model_config_yaml(self):
|
||||
return self._model_config_yaml
|
||||
|
||||
@property
|
||||
def model_weights_ckpt(self):
|
||||
return self._model_weights_ckpt
|
||||
|
||||
@property
|
||||
def model_restore_path(self):
|
||||
restore_path = self._all_model_restore_paths[-1] if len(self._all_model_restore_paths) > 0 else None
|
||||
return restore_path
|
||||
|
||||
@model_restore_path.setter
|
||||
def model_restore_path(self, path):
|
||||
with self.__lock:
|
||||
self._model_restore_path = path
|
||||
self._all_model_restore_paths.append(path)
|
||||
|
||||
def register_model_guid(self, guid: str, restoration_path: Optional[str] = None):
|
||||
# Maps a guid to its restore path (None or last absolute path)
|
||||
with self.__lock:
|
||||
if guid in self._model_guid_map:
|
||||
idx = self._model_guid_map[guid].gidx
|
||||
else:
|
||||
idx = len(self._model_guid_map)
|
||||
self._model_guid_map[guid] = ModelMetadataRegistry(guid, idx, restoration_path=restoration_path)
|
||||
|
||||
def reset_model_guid_registry(self):
|
||||
# Reset the guid mapping
|
||||
with self.__lock:
|
||||
self._model_guid_map.clear()
|
||||
|
||||
def get_model_metadata_from_guid(self, guid) -> ModelMetadataRegistry:
|
||||
# Returns the global model idx and restoration path
|
||||
metadata = self._model_guid_map[guid]
|
||||
return metadata
|
||||
@@ -0,0 +1,132 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from argparse import ArgumentParser
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
|
||||
def add_optimizer_args(
|
||||
parent_parser: ArgumentParser,
|
||||
optimizer: str = 'adam',
|
||||
default_lr: float = None,
|
||||
default_opt_args: Optional[Union[Dict[str, Any], List[str]]] = None,
|
||||
) -> ArgumentParser:
|
||||
"""Extends existing argparse with support for optimizers.
|
||||
|
||||
# Example of adding optimizer args to command line :
|
||||
python train_script.py ... --optimizer "novograd" --lr 0.01 \
|
||||
--opt_args betas=0.95,0.5 weight_decay=0.001
|
||||
|
||||
Args:
|
||||
parent_parser (ArgumentParser): Custom CLI parser that will be extended.
|
||||
optimizer (str): Default optimizer required.
|
||||
default_lr (float): Default learning rate that should be overriden during training.
|
||||
default_opt_args (list(str)): List of overriding arguments for the instantiated optimizer.
|
||||
|
||||
Returns:
|
||||
ArgumentParser: Parser extended by Optimizers arguments.
|
||||
"""
|
||||
if default_opt_args is None:
|
||||
default_opt_args = []
|
||||
|
||||
parser = ArgumentParser(parents=[parent_parser], add_help=True, conflict_handler='resolve')
|
||||
|
||||
parser.add_argument('--optimizer', type=str, default=optimizer, help='Name of the optimizer. Defaults to Adam.')
|
||||
parser.add_argument('--lr', type=float, default=default_lr, help='Learning rate of the optimizer.')
|
||||
parser.add_argument(
|
||||
'--opt_args',
|
||||
default=default_opt_args,
|
||||
nargs='+',
|
||||
type=str,
|
||||
help='Overriding arguments for the optimizer. \n Must follow the pattern : \n name=value separated by spaces.'
|
||||
'Example: --opt_args weight_decay=0.001 eps=1e-8 betas=0.9,0.999',
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def add_scheduler_args(parent_parser: ArgumentParser) -> ArgumentParser:
|
||||
"""Extends existing argparse with default LR scheduler args.
|
||||
|
||||
Args:
|
||||
parent_parser (ArgumentParser): Custom CLI parser that will be extended.
|
||||
|
||||
Returns:
|
||||
ArgumentParser: Parser extended by LR Scheduler arguments.
|
||||
"""
|
||||
parser = ArgumentParser(parents=[parent_parser], add_help=False, conflict_handler='resolve')
|
||||
parser.add_argument("--warmup_steps", type=int, required=False, default=None, help="Number of warmup steps")
|
||||
parser.add_argument(
|
||||
"--warmup_ratio",
|
||||
type=float,
|
||||
required=False,
|
||||
default=None,
|
||||
help="Number of warmup steps as a percentage of total training steps",
|
||||
)
|
||||
parser.add_argument("--hold_steps", type=int, required=False, default=None, help="Number of hold LR steps")
|
||||
parser.add_argument(
|
||||
"--hold_ratio",
|
||||
type=float,
|
||||
required=False,
|
||||
default=None,
|
||||
help="Number of hold LR steps as a percentage of total training steps",
|
||||
)
|
||||
parser.add_argument("--min_lr", type=float, required=False, default=0.0, help="Minimum learning rate")
|
||||
parser.add_argument(
|
||||
"--last_epoch", type=int, required=False, default=-1, help="Last epoch id. -1 indicates training from scratch"
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def add_asr_args(parent_parser: ArgumentParser) -> ArgumentParser:
|
||||
"""Extends existing argparse with default ASR collection args.
|
||||
|
||||
Args:
|
||||
parent_parser (ArgumentParser): Custom CLI parser that will be extended.
|
||||
|
||||
Returns:
|
||||
ArgumentParser: Parser extended by NeMo ASR Collection arguments.
|
||||
"""
|
||||
parser = ArgumentParser(parents=[parent_parser], add_help=False, conflict_handler='resolve')
|
||||
parser.add_argument("--asr_model", type=str, required=True, default="bad_quartznet15x5.yaml", help="")
|
||||
parser.add_argument("--train_dataset", type=str, required=True, default=None, help="training dataset path")
|
||||
parser.add_argument("--eval_dataset", type=str, required=True, help="evaluation dataset path")
|
||||
return parser
|
||||
|
||||
|
||||
def add_nlp_args(parent_parser: ArgumentParser) -> ArgumentParser:
|
||||
"""Extends existing argparse with default NLP collection args.
|
||||
|
||||
Args:
|
||||
parent_parser (ArgumentParser): Custom CLI parser that will be extended.
|
||||
|
||||
Returns:
|
||||
ArgumentParser: Parser extended by NeMo NLP Collection arguments.
|
||||
"""
|
||||
parser = ArgumentParser(parents=[parent_parser], add_help=False, conflict_handler='resolve')
|
||||
parser.add_argument(
|
||||
"--data_dir", type=str, required=False, help="data directory to training or/and evaluation dataset"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--config_file", type=str, required=False, default=None, help="Huggingface model configuration file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pretrained_model_name", default='bert-base-uncased', type=str, required=False, help="pretrained model name"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tokenizer_name", default='nemobert', type=str, choices=['sentencepiece', 'nemobert'], help="Tokenizer type"
|
||||
)
|
||||
parser.add_argument("--tokenizer_model", default=None, type=str, help="Tokenizer file for sentence piece")
|
||||
parser.add_argument("--do_lower_case", action='store_true', required=False, help="lower case data")
|
||||
return parser
|
||||
@@ -0,0 +1,82 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from time import sleep
|
||||
|
||||
import wget
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
|
||||
def maybe_download_from_cloud(url, filename, subfolder=None, cache_dir=None, refresh_cache=False) -> str:
|
||||
"""
|
||||
Helper function to download pre-trained weights from the cloud
|
||||
Args:
|
||||
url: (str) URL of storage
|
||||
filename: (str) what to download. The request will be issued to url/filename
|
||||
subfolder: (str) subfolder within cache_dir. The file will be stored in cache_dir/subfolder. Subfolder can
|
||||
be empty
|
||||
cache_dir: (str) a cache directory where to download. If not present, this function will attempt to create it.
|
||||
If None (default), then it will be $HOME/.cache/torch/NeMo
|
||||
refresh_cache: (bool) if True and cached file is present, it will delete it and re-fetch
|
||||
|
||||
Returns:
|
||||
If successful - absolute local path to the downloaded file
|
||||
else - empty string
|
||||
"""
|
||||
# try:
|
||||
if cache_dir is None:
|
||||
cache_location = Path.joinpath(Path.home(), ".cache/torch/NeMo")
|
||||
else:
|
||||
cache_location = cache_dir
|
||||
if subfolder is not None:
|
||||
destination = Path.joinpath(cache_location, subfolder)
|
||||
else:
|
||||
destination = cache_location
|
||||
|
||||
if not os.path.exists(destination):
|
||||
os.makedirs(destination, exist_ok=True)
|
||||
|
||||
destination_file = Path.joinpath(destination, filename)
|
||||
|
||||
if os.path.exists(destination_file):
|
||||
logging.info(f"Found existing object {destination_file}.")
|
||||
if refresh_cache:
|
||||
logging.info("Asked to refresh the cache.")
|
||||
logging.info(f"Deleting file: {destination_file}")
|
||||
os.remove(destination_file)
|
||||
else:
|
||||
logging.info(f"Re-using file from: {destination_file}")
|
||||
return str(destination_file)
|
||||
# download file
|
||||
wget_uri = url + filename
|
||||
logging.info(f"Downloading from: {wget_uri} to {str(destination_file)}")
|
||||
# NGC links do not work everytime so we try and wait
|
||||
i = 0
|
||||
max_attempts = 3
|
||||
while i < max_attempts:
|
||||
i += 1
|
||||
try:
|
||||
wget.download(wget_uri, str(destination_file))
|
||||
if os.path.exists(destination_file):
|
||||
return destination_file
|
||||
else:
|
||||
return ""
|
||||
except:
|
||||
logging.info(f"Download from cloud failed. Attempt {i} of {max_attempts}")
|
||||
sleep(0.05)
|
||||
continue
|
||||
raise ValueError("Not able to download url right now, please try again.")
|
||||
@@ -0,0 +1,258 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import copy
|
||||
import inspect
|
||||
from dataclasses import is_dataclass
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from omegaconf import DictConfig, OmegaConf, open_dict
|
||||
|
||||
from nemo.core.config.modelPT import NemoConfig
|
||||
from nemo.utils import logging
|
||||
|
||||
|
||||
def update_model_config(model_cls: NemoConfig, update_cfg: DictConfig, drop_missing_subconfigs: bool = True):
|
||||
"""
|
||||
Helper class that updates the default values of a ModelPT config class with the values
|
||||
in a DictConfig that mirrors the structure of the config class.
|
||||
|
||||
Assumes the `update_cfg` is a DictConfig (either generated manually, via hydra or instantiated via yaml/model.cfg).
|
||||
This update_cfg is then used to override the default values preset inside the ModelPT config class.
|
||||
|
||||
If `drop_missing_subconfigs` is set, the certain sub-configs of the ModelPT config class will be removed, iff
|
||||
they are not found in the mirrored `update_cfg`. The following sub-configs are subject to potential removal:
|
||||
- `train_ds`
|
||||
- `validation_ds`
|
||||
- `test_ds`
|
||||
- `optim` + nested `sched`.
|
||||
|
||||
Args:
|
||||
model_cls: A subclass of NemoConfig, that details in entirety all of the parameters that constitute
|
||||
the NeMo Model.
|
||||
|
||||
update_cfg: A DictConfig that mirrors the structure of the NemoConfig data class. Used to update the
|
||||
default values of the config class.
|
||||
|
||||
drop_missing_subconfigs: Bool which determins whether to drop certain sub-configs from the NemoConfig
|
||||
class, if the corresponding sub-config is missing from `update_cfg`.
|
||||
|
||||
Returns:
|
||||
A DictConfig with updated values that can be used to instantiate the NeMo Model along with supporting
|
||||
infrastructure.
|
||||
"""
|
||||
if not (is_dataclass(model_cls) or isinstance(model_cls, DictConfig)):
|
||||
raise ValueError("`model_cfg` must be a dataclass or a structured OmegaConf object")
|
||||
|
||||
if not isinstance(update_cfg, DictConfig):
|
||||
update_cfg = OmegaConf.create(update_cfg)
|
||||
|
||||
if is_dataclass(model_cls):
|
||||
model_cls = OmegaConf.structured(model_cls)
|
||||
|
||||
# Update optional configs
|
||||
model_cls = _update_subconfig(
|
||||
model_cls, update_cfg, subconfig_key='train_ds', drop_missing_subconfigs=drop_missing_subconfigs
|
||||
)
|
||||
model_cls = _update_subconfig(
|
||||
model_cls, update_cfg, subconfig_key='validation_ds', drop_missing_subconfigs=drop_missing_subconfigs
|
||||
)
|
||||
model_cls = _update_subconfig(
|
||||
model_cls, update_cfg, subconfig_key='test_ds', drop_missing_subconfigs=drop_missing_subconfigs
|
||||
)
|
||||
model_cls = _update_subconfig(
|
||||
model_cls, update_cfg, subconfig_key='optim', drop_missing_subconfigs=drop_missing_subconfigs
|
||||
)
|
||||
|
||||
# Add optim and sched additional keys to model cls
|
||||
model_cls = _add_subconfig_keys(model_cls, update_cfg, subconfig_key='optim')
|
||||
|
||||
# Perform full merge of model config class and update config
|
||||
# Remove ModelPT artifact `target`
|
||||
if 'target' in update_cfg.model:
|
||||
# Assume artifact from ModelPT and pop
|
||||
if 'target' not in model_cls.model:
|
||||
with open_dict(update_cfg.model):
|
||||
update_cfg.model.pop('target')
|
||||
|
||||
model_cfg = OmegaConf.merge(model_cls, update_cfg)
|
||||
|
||||
return model_cfg
|
||||
|
||||
|
||||
def _update_subconfig(
|
||||
model_cfg: DictConfig, update_cfg: DictConfig, subconfig_key: str, drop_missing_subconfigs: bool
|
||||
):
|
||||
"""
|
||||
Updates the NemoConfig DictConfig such that:
|
||||
1) If the sub-config key exists in the `update_cfg`, but does not exist in ModelPT config:
|
||||
- Add the sub-config from update_cfg to ModelPT config
|
||||
|
||||
2) If the sub-config key does not exist in `update_cfg`, but exists in ModelPT config:
|
||||
- Remove the sub-config from the ModelPT config; iff the `drop_missing_subconfigs` flag is set.
|
||||
|
||||
Args:
|
||||
model_cfg: A DictConfig instantiated from the NemoConfig subclass.
|
||||
update_cfg: A DictConfig that mirrors the structure of `model_cfg`, used to update its default values.
|
||||
subconfig_key: A str key used to check and update the sub-config.
|
||||
drop_missing_subconfigs: A bool flag, whether to allow deletion of the NemoConfig sub-config,
|
||||
if its mirror sub-config does not exist in the `update_cfg`.
|
||||
|
||||
Returns:
|
||||
The updated DictConfig for the NemoConfig
|
||||
"""
|
||||
with open_dict(model_cfg.model):
|
||||
# If update config has the key, but model cfg doesnt have the key
|
||||
# Add the update cfg subconfig to the model cfg
|
||||
if subconfig_key in update_cfg.model and subconfig_key not in model_cfg.model:
|
||||
model_cfg.model[subconfig_key] = update_cfg.model[subconfig_key]
|
||||
|
||||
# If update config does not the key, but model cfg has the key
|
||||
# Remove the model cfg subconfig in order to match layout of update cfg
|
||||
if subconfig_key not in update_cfg.model and subconfig_key in model_cfg.model:
|
||||
if drop_missing_subconfigs:
|
||||
model_cfg.model.pop(subconfig_key)
|
||||
|
||||
return model_cfg
|
||||
|
||||
|
||||
def _add_subconfig_keys(model_cfg: DictConfig, update_cfg: DictConfig, subconfig_key: str):
|
||||
"""
|
||||
For certain sub-configs, the default values specified by the NemoConfig class is insufficient.
|
||||
In order to support every potential value in the merge between the `update_cfg`, it would require
|
||||
explicit definition of all possible cases.
|
||||
|
||||
An example of such a case is Optimizers, and their equivalent Schedulers. All optimizers share a few basic
|
||||
details - such as name and lr, but almost all require additional parameters - such as weight decay.
|
||||
It is impractical to create a config for every single optimizer + every single scheduler combination.
|
||||
|
||||
In such a case, we perform a dual merge. The Optim and Sched Dataclass contain the bare minimum essential
|
||||
components. The extra values are provided via update_cfg.
|
||||
|
||||
In order to enable the merge, we first need to update the update sub-config to incorporate the keys,
|
||||
with dummy temporary values (merge update config with model config). This is done on a copy of the
|
||||
update sub-config, as the actual override values might be overriden by the NemoConfig defaults.
|
||||
|
||||
Then we perform a merge of this temporary sub-config with the actual override config in a later step
|
||||
(merge model_cfg with original update_cfg, done outside this function).
|
||||
|
||||
Args:
|
||||
model_cfg: A DictConfig instantiated from the NemoConfig subclass.
|
||||
update_cfg: A DictConfig that mirrors the structure of `model_cfg`, used to update its default values.
|
||||
subconfig_key: A str key used to check and update the sub-config.
|
||||
|
||||
Returns:
|
||||
A ModelPT DictConfig with additional keys added to the sub-config.
|
||||
"""
|
||||
with open_dict(model_cfg.model):
|
||||
# Create copy of original model sub config
|
||||
if subconfig_key in update_cfg.model:
|
||||
if subconfig_key not in model_cfg.model:
|
||||
# create the key as a placeholder
|
||||
model_cfg.model[subconfig_key] = None
|
||||
|
||||
subconfig = copy.deepcopy(model_cfg.model[subconfig_key])
|
||||
update_subconfig = copy.deepcopy(update_cfg.model[subconfig_key])
|
||||
|
||||
# Add the keys and update temporary values, will be updated during full merge
|
||||
subconfig = OmegaConf.merge(update_subconfig, subconfig)
|
||||
# Update sub config
|
||||
model_cfg.model[subconfig_key] = subconfig
|
||||
|
||||
return model_cfg
|
||||
|
||||
|
||||
def assert_dataclass_signature_match(
|
||||
cls: 'class_type',
|
||||
datacls: 'dataclass',
|
||||
ignore_args: Optional[List[str]] = None,
|
||||
remap_args: Optional[Dict[str, str]] = None,
|
||||
):
|
||||
"""
|
||||
Analyses the signature of a provided class and its respective data class,
|
||||
asserting that the dataclass signature matches the class __init__ signature.
|
||||
|
||||
Note:
|
||||
This is not a value based check. This function only checks if all argument
|
||||
names exist on both class and dataclass and logs mismatches.
|
||||
|
||||
Args:
|
||||
cls: Any class type - but not an instance of a class. Pass type(x) where x is an instance
|
||||
if class type is not easily available.
|
||||
datacls: A corresponding dataclass for the above class.
|
||||
ignore_args: (Optional) A list of string argument names which are forcibly ignored,
|
||||
even if mismatched in the signature. Useful when a dataclass is a superset of the
|
||||
arguments of a class.
|
||||
remap_args: (Optional) A dictionary, mapping an argument name that exists (in either the
|
||||
class or its dataclass), to another name. Useful when argument names are mismatched between
|
||||
a class and its dataclass due to indirect instantiation via a helper method.
|
||||
|
||||
Returns:
|
||||
A tuple containing information about the analysis:
|
||||
1) A bool value which is True if the signatures matched exactly / after ignoring values.
|
||||
False otherwise.
|
||||
2) A set of arguments names that exist in the class, but *do not* exist in the dataclass.
|
||||
If exact signature match occurs, this will be None instead.
|
||||
3) A set of argument names that exist in the data class, but *do not* exist in the class itself.
|
||||
If exact signature match occurs, this will be None instead.
|
||||
"""
|
||||
class_sig = inspect.signature(cls.__init__)
|
||||
|
||||
class_params = dict(**class_sig.parameters)
|
||||
class_params.pop('self')
|
||||
|
||||
dataclass_sig = inspect.signature(datacls)
|
||||
|
||||
dataclass_params = dict(**dataclass_sig.parameters)
|
||||
dataclass_params.pop("_target_", None)
|
||||
|
||||
class_params = set(class_params.keys())
|
||||
dataclass_params = set(dataclass_params.keys())
|
||||
|
||||
if remap_args is not None:
|
||||
for original_arg, new_arg in remap_args.items():
|
||||
if original_arg in class_params:
|
||||
class_params.remove(original_arg)
|
||||
class_params.add(new_arg)
|
||||
logging.info(f"Remapped {original_arg} -> {new_arg} in {cls.__name__}")
|
||||
|
||||
if original_arg in dataclass_params:
|
||||
dataclass_params.remove(original_arg)
|
||||
dataclass_params.add(new_arg)
|
||||
logging.info(f"Remapped {original_arg} -> {new_arg} in {datacls.__name__}")
|
||||
|
||||
if ignore_args is not None:
|
||||
ignore_args = set(ignore_args)
|
||||
|
||||
class_params = class_params - ignore_args
|
||||
dataclass_params = dataclass_params - ignore_args
|
||||
logging.info(f"Removing ignored arguments - {ignore_args}")
|
||||
|
||||
intersection = set.intersection(class_params, dataclass_params)
|
||||
subset_cls = class_params - intersection
|
||||
subset_datacls = dataclass_params - intersection
|
||||
|
||||
if (len(class_params) != len(dataclass_params)) or len(subset_cls) > 0 or len(subset_datacls) > 0:
|
||||
logging.error(f"Class {cls.__name__} arguments do not match " f"Dataclass {datacls.__name__}!")
|
||||
|
||||
if len(subset_cls) > 0:
|
||||
logging.error(f"Class {cls.__name__} has additional arguments :\n" f"{subset_cls}")
|
||||
|
||||
if len(subset_datacls):
|
||||
logging.error(f"Dataclass {datacls.__name__} has additional arguments :\n{subset_datacls}")
|
||||
|
||||
return False, subset_cls, subset_datacls
|
||||
|
||||
else:
|
||||
return True, None, None
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
from nemo.utils.decorators.deprecated import deprecated
|
||||
from nemo.utils.decorators.experimental import experimental
|
||||
from nemo.utils.decorators.port_docs import add_port_docs
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
__all__ = [
|
||||
'deprecated',
|
||||
]
|
||||
|
||||
import functools
|
||||
|
||||
import wrapt
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
# Remember which deprecation warnings have been printed already.
|
||||
_PRINTED_WARNING = {}
|
||||
|
||||
|
||||
def deprecated(wrapped=None, version=None, explanation=None):
|
||||
""" Decorator class used for indicating that a function is deprecated and going to be removed.
|
||||
Tracks down which functions printed the warning and will print it only once per function.
|
||||
"""
|
||||
|
||||
if wrapped is None:
|
||||
return functools.partial(deprecated, version=version, explanation=explanation)
|
||||
|
||||
@wrapt.decorator
|
||||
def wrapper(wrapped, instance, args, kwargs):
|
||||
"""
|
||||
Method prints the adequate warning (only once per function) when
|
||||
required and calls the function func, passing the original arguments,
|
||||
i.e. version and explanation.
|
||||
|
||||
Args:
|
||||
version: Version in which the function will be removed (optional)
|
||||
explanation: Additional explanation (optional), e.g. use method ``blabla instead``.
|
||||
"""
|
||||
|
||||
# Check if we already warned about that function.
|
||||
if wrapped.__name__ not in _PRINTED_WARNING.keys():
|
||||
# Add to list so we won't print it again.
|
||||
_PRINTED_WARNING[wrapped.__name__] = True
|
||||
|
||||
# Prepare the warning message.
|
||||
msg = "Function ``{}`` is deprecated.".format(wrapped.__name__)
|
||||
|
||||
# Optionally, add version and alternative.
|
||||
if version is not None:
|
||||
msg = msg + " It is going to be removed in "
|
||||
msg = msg + "the {} version.".format(version)
|
||||
|
||||
if explanation is not None:
|
||||
msg = msg + " " + explanation
|
||||
|
||||
# Display the deprecated warning.
|
||||
logging.warning(msg)
|
||||
|
||||
# Call the function.
|
||||
return wrapped(*args, **kwargs)
|
||||
|
||||
return wrapper(wrapped)
|
||||
@@ -0,0 +1,33 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
__all__ = ['experimental']
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
|
||||
def experimental(cls):
|
||||
""" Decorator which indicates that module is experimental.
|
||||
Use it to mark experimental or research modules.
|
||||
"""
|
||||
|
||||
def wrapped(cls):
|
||||
logging.warning(
|
||||
f'Module {cls} is experimental, not ready for production and is not fully supported. Use at your own risk.'
|
||||
)
|
||||
|
||||
return cls
|
||||
|
||||
return wrapped(cls=cls)
|
||||
@@ -0,0 +1,90 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
# The "add_port_docs" decorator is needed to nicely generate neural types in Sphynx for input and output ports
|
||||
|
||||
__all__ = [
|
||||
'add_port_docs',
|
||||
]
|
||||
|
||||
import functools
|
||||
import sys
|
||||
|
||||
import wrapt
|
||||
|
||||
|
||||
def _normalize_docstring(docstring):
|
||||
"""Normalizes the docstring.
|
||||
Replaces tabs with spaces, removes leading and trailing blanks lines, and
|
||||
removes any indentation.
|
||||
Copied from PEP-257:
|
||||
https://www.python.org/dev/peps/pep-0257/#handling-docstring-indentation
|
||||
Args:
|
||||
docstring: the docstring to normalize
|
||||
Returns:
|
||||
The normalized docstring
|
||||
"""
|
||||
if not docstring:
|
||||
return ''
|
||||
# Convert tabs to spaces (following the normal Python rules)
|
||||
# and split into a list of lines:
|
||||
lines = docstring.expandtabs().splitlines()
|
||||
# Determine minimum indentation (first line doesn't count):
|
||||
# (we use sys.maxsize because sys.maxint doesn't exist in Python 3)
|
||||
indent = sys.maxsize
|
||||
for line in lines[1:]:
|
||||
stripped = line.lstrip()
|
||||
if stripped:
|
||||
indent = min(indent, len(line) - len(stripped))
|
||||
# Remove indentation (first line is special):
|
||||
trimmed = [lines[0].strip()]
|
||||
if indent < sys.maxsize:
|
||||
for line in lines[1:]:
|
||||
trimmed.append(line[indent:].rstrip())
|
||||
# Strip off trailing and leading blank lines:
|
||||
while trimmed and not trimmed[-1]:
|
||||
trimmed.pop()
|
||||
while trimmed and not trimmed[0]:
|
||||
trimmed.pop(0)
|
||||
# Return a single string:
|
||||
return '\n'.join(trimmed)
|
||||
|
||||
|
||||
def add_port_docs(wrapped=None, instance=None, value=''):
|
||||
if wrapped is None:
|
||||
return functools.partial(add_port_docs, value=value)
|
||||
|
||||
@wrapt.decorator
|
||||
def wrapper(wrapped, instance=None, args=None, kwargs=None):
|
||||
return wrapped(*args, **kwargs)
|
||||
|
||||
decorated = wrapper(wrapped)
|
||||
try:
|
||||
port_2_ntype = decorated(instance)
|
||||
except:
|
||||
port_2_ntype = None
|
||||
|
||||
port_description = ""
|
||||
if port_2_ntype is not None:
|
||||
for port, ntype in port_2_ntype.items():
|
||||
port_description += "* *" + port + "* : " + str(ntype)
|
||||
port_description += "\n\n"
|
||||
|
||||
__doc__ = _normalize_docstring(wrapped.__doc__) + '\n\n' + str(port_description)
|
||||
__doc__ = _normalize_docstring(__doc__)
|
||||
|
||||
wrapt.FunctionWrapper.__setattr__(decorated, "__doc__", __doc__)
|
||||
|
||||
return decorated
|
||||
@@ -0,0 +1,207 @@
|
||||
# The MIT Licence (MIT)
|
||||
#
|
||||
# Copyright (c) 2016 YunoJuno Ltd
|
||||
#
|
||||
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
# of this software and associated documentation files (the "Software"), to deal
|
||||
# in the Software without restriction, including without limitation the rights
|
||||
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
# copies of the Software, and to permit persons to whom the Software is
|
||||
# furnished to do so, subject to the following conditions:
|
||||
#
|
||||
# The above copyright notice and this permission notice shall be included in all
|
||||
# copies or substantial portions of the Software.
|
||||
#
|
||||
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
# SOFTWARE.
|
||||
#
|
||||
# Vendored dependency from : https://github.com/yunojuno/python-env-utils/blob/master/env_utils/utils.py
|
||||
#
|
||||
#
|
||||
# Modified by NVIDIA
|
||||
#
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import decimal
|
||||
import json
|
||||
import os
|
||||
|
||||
from dateutil import parser
|
||||
|
||||
__all__ = [
|
||||
"get_env",
|
||||
"get_envbool",
|
||||
"get_envint",
|
||||
"get_envfloat",
|
||||
"get_envdecimal",
|
||||
"get_envdate",
|
||||
"get_envdatetime",
|
||||
"get_envlist",
|
||||
"get_envdict",
|
||||
"CoercionError",
|
||||
"RequiredSettingMissingError",
|
||||
]
|
||||
|
||||
|
||||
class CoercionError(Exception):
|
||||
"""Custom error raised when a value cannot be coerced."""
|
||||
|
||||
def __init__(self, key, value, func):
|
||||
msg = "Unable to coerce '{}={}' using {}.".format(key, value, func.__name__)
|
||||
super(CoercionError, self).__init__(msg)
|
||||
|
||||
|
||||
class RequiredSettingMissingError(Exception):
|
||||
"""Custom error raised when a required env var is missing."""
|
||||
|
||||
def __init__(self, key):
|
||||
msg = "Required env var '{}' is missing.".format(key)
|
||||
super(RequiredSettingMissingError, self).__init__(msg)
|
||||
|
||||
|
||||
def _get_env(key, default=None, coerce=lambda x: x, required=False):
|
||||
"""
|
||||
Return env var coerced into a type other than string.
|
||||
This function extends the standard os.getenv function to enable
|
||||
the coercion of values into data types other than string (all env
|
||||
vars are strings by default).
|
||||
Args:
|
||||
key: string, the name of the env var to look up
|
||||
Kwargs:
|
||||
default: the default value to return if the env var does not exist. NB the
|
||||
default value is **not** coerced, and is assumed to be of the correct type.
|
||||
coerce: a function that is used to coerce the value returned into
|
||||
another type
|
||||
required: bool, if True, then a RequiredSettingMissingError error is raised
|
||||
if the env var does not exist.
|
||||
Returns the env var, passed through the coerce function
|
||||
"""
|
||||
try:
|
||||
value = os.environ[key]
|
||||
except KeyError:
|
||||
if required is True:
|
||||
raise RequiredSettingMissingError(key)
|
||||
else:
|
||||
return default
|
||||
|
||||
try:
|
||||
return coerce(value)
|
||||
except Exception:
|
||||
raise CoercionError(key, value, coerce)
|
||||
|
||||
|
||||
# standard type coercion functions
|
||||
def _bool(value):
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
|
||||
return not (value is None or value.lower() in ("false", "0", "no", "n", "f", "none"))
|
||||
|
||||
|
||||
def _int(value):
|
||||
return int(value)
|
||||
|
||||
|
||||
def _float(value):
|
||||
return float(value)
|
||||
|
||||
|
||||
def _decimal(value):
|
||||
return decimal.Decimal(value)
|
||||
|
||||
|
||||
def _dict(value):
|
||||
return json.loads(value)
|
||||
|
||||
|
||||
def _datetime(value):
|
||||
return parser.parse(value)
|
||||
|
||||
|
||||
def _date(value):
|
||||
return parser.parse(value).date()
|
||||
|
||||
|
||||
def get_env(key, *default, **kwargs):
|
||||
"""
|
||||
Return env var.
|
||||
This is the parent function of all other get_foo functions,
|
||||
and is responsible for unpacking args/kwargs into the values
|
||||
that _get_env expects (it is the root function that actually
|
||||
interacts with environ).
|
||||
Args:
|
||||
key: string, the env var name to look up.
|
||||
default: (optional) the value to use if the env var does not
|
||||
exist. If this value is not supplied, then the env var is
|
||||
considered to be required, and a RequiredSettingMissingError
|
||||
error will be raised if it does not exist.
|
||||
Kwargs:
|
||||
coerce: a func that may be supplied to coerce the value into
|
||||
something else. This is used by the default get_foo functions
|
||||
to cast strings to builtin types, but could be a function that
|
||||
returns a custom class.
|
||||
Returns the env var, coerced if required, and a default if supplied.
|
||||
"""
|
||||
assert len(default) in (0, 1), "Too many args supplied."
|
||||
func = kwargs.get('coerce', lambda x: x)
|
||||
required = len(default) == 0
|
||||
default = default[0] if not required else None
|
||||
return _get_env(key, default=default, coerce=func, required=required)
|
||||
|
||||
|
||||
def get_envbool(key, *default):
|
||||
"""Return env var cast as boolean."""
|
||||
return get_env(key, *default, coerce=_bool)
|
||||
|
||||
|
||||
def get_envint(key, *default):
|
||||
"""Return env var cast as integer."""
|
||||
return get_env(key, *default, coerce=_int)
|
||||
|
||||
|
||||
def get_envfloat(key, *default):
|
||||
"""Return env var cast as float."""
|
||||
return get_env(key, *default, coerce=_float)
|
||||
|
||||
|
||||
def get_envdecimal(key, *default):
|
||||
"""Return env var cast as Decimal."""
|
||||
return get_env(key, *default, coerce=_decimal)
|
||||
|
||||
|
||||
def get_envdate(key, *default):
|
||||
"""Return env var as a date."""
|
||||
return get_env(key, *default, coerce=_date)
|
||||
|
||||
|
||||
def get_envdatetime(key, *default):
|
||||
"""Return env var as a datetime."""
|
||||
return get_env(key, *default, coerce=_datetime)
|
||||
|
||||
|
||||
def get_envlist(key, *default, **kwargs):
|
||||
"""Return env var as a list."""
|
||||
separator = kwargs.get('separator', ' ')
|
||||
return get_env(key, *default, coerce=lambda x: x.split(separator))
|
||||
|
||||
|
||||
def get_envdict(key, *default):
|
||||
"""Return env var as a dict."""
|
||||
return get_env(key, *default, coerce=_dict)
|
||||
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
class NeMoBaseException(Exception):
|
||||
""" NeMo Base Exception. All exceptions created in NeMo should inherit from this class"""
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,715 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from shutil import copy, move
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import torch
|
||||
from hydra.core.hydra_config import HydraConfig
|
||||
from hydra.utils import get_original_cwd
|
||||
from omegaconf import DictConfig, OmegaConf, open_dict
|
||||
from pytorch_lightning.callbacks import ModelCheckpoint
|
||||
from pytorch_lightning.loggers import LoggerCollection as _LoggerCollection
|
||||
from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
|
||||
from pytorch_lightning.utilities import rank_zero_only
|
||||
|
||||
from nemo.constants import NEMO_ENV_VARNAME_VERSION
|
||||
from nemo.utils import app_state, logging
|
||||
from nemo.utils.app_state import AppState
|
||||
from nemo.utils.exceptions import NeMoBaseException
|
||||
from nemo.utils.get_rank import is_global_rank_zero
|
||||
from nemo.utils.lightning_logger_patch import add_filehandlers_to_pl_logger
|
||||
|
||||
|
||||
class NotFoundError(NeMoBaseException):
|
||||
""" Raised when a file or folder is not found"""
|
||||
|
||||
|
||||
class LoggerMisconfigurationError(NeMoBaseException):
|
||||
""" Raised when a mismatch between trainer.logger and exp_manager occurs"""
|
||||
|
||||
def __init__(self, message):
|
||||
message = (
|
||||
message
|
||||
+ " You can disable lighning's trainer from creating a logger by passing logger=False to its constructor."
|
||||
)
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class CheckpointMisconfigurationError(NeMoBaseException):
|
||||
""" Raised when a mismatch between trainer.callbacks and exp_manager occurs"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class CallbackParams:
|
||||
filepath: Optional[str] = None # Deprecated
|
||||
dirpath: Optional[str] = None # If None, exp_manager will attempt to handle the filepath
|
||||
filename: Optional[str] = None # If None, exp_manager will attempt to handle the filepath
|
||||
monitor: Optional[str] = "val_loss"
|
||||
verbose: Optional[bool] = True
|
||||
save_last: Optional[bool] = True
|
||||
save_top_k: Optional[int] = 3
|
||||
save_weights_only: Optional[bool] = False
|
||||
mode: Optional[str] = "min"
|
||||
period: Optional[int] = 1
|
||||
prefix: Optional[str] = None # If None, exp_manager will attempt to handle the filepath
|
||||
postfix: str = ".nemo"
|
||||
save_best_model: bool = False
|
||||
always_save_nemo: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExpManagerConfig:
|
||||
# Log dir creation parameters
|
||||
explicit_log_dir: Optional[str] = None
|
||||
exp_dir: Optional[str] = None
|
||||
name: Optional[str] = None
|
||||
version: Optional[str] = None
|
||||
use_datetime_version: Optional[bool] = True
|
||||
resume_if_exists: Optional[bool] = False
|
||||
resume_past_end: Optional[bool] = False
|
||||
resume_ignore_no_checkpoint: Optional[bool] = False
|
||||
# Logging parameters
|
||||
create_tensorboard_logger: Optional[bool] = True
|
||||
summary_writer_kwargs: Optional[Dict[Any, Any]] = None
|
||||
create_wandb_logger: Optional[bool] = False
|
||||
wandb_logger_kwargs: Optional[Dict[Any, Any]] = None
|
||||
# Checkpointing parameters
|
||||
create_checkpoint_callback: Optional[bool] = True
|
||||
checkpoint_callback_params: Optional[CallbackParams] = CallbackParams()
|
||||
# Additional exp_manager arguments
|
||||
files_to_copy: Optional[List[str]] = None
|
||||
|
||||
|
||||
def exp_manager(trainer: 'pytorch_lightning.Trainer', cfg: Optional[Union[DictConfig, Dict]] = None) -> Path:
|
||||
"""
|
||||
exp_manager is a helper function used to manage folders for experiments. It follows the pytorch lightning paradigm
|
||||
of exp_dir/model_or_experiment_name/version. If the lightning trainer has a logger, exp_manager will get exp_dir,
|
||||
name, and version from the logger. Otherwise it will use the exp_dir and name arguments to create the logging
|
||||
directory. exp_manager also allows for explicit folder creation via explicit_log_dir.
|
||||
|
||||
The version can be a datetime string or an integer. Datestime version can be disabled if use_datetime_version is set
|
||||
to False. It optionally creates TensorBoardLogger, WandBLogger, ModelCheckpoint objects from pytorch lightning.
|
||||
It copies sys.argv, and git information if available to the logging directory. It creates a log file for each
|
||||
process to log their output into.
|
||||
|
||||
exp_manager additionally has a resume feature (resume_if_exists) which can be used to continuing training from
|
||||
the constructed log_dir. When you need to continue the training repeatedly (like on a cluster which you need
|
||||
multiple consecutive jobs), you need to avoid creating the version folders. Therefore from v1.0.0, when
|
||||
resume_if_exists is set to True, creating the version folders is ignored.
|
||||
|
||||
Args:
|
||||
trainer (pytorch_lightning.Trainer): The lightning trainer.
|
||||
cfg (DictConfig, dict): Can have the following keys:
|
||||
- explicit_log_dir (str, Path): Can be used to override exp_dir/name/version folder creation. Defaults to
|
||||
None, which will use exp_dir, name, and version to construct the logging directory.
|
||||
- exp_dir (str, Path): The base directory to create the logging directory. Defaults to None, which logs to
|
||||
./nemo_experiments.
|
||||
- name (str): The name of the experiment. Defaults to None which turns into "default" via name = name or
|
||||
"default".
|
||||
- version (str): The version of the experiment. Defaults to None which uses either a datetime string or
|
||||
lightning's TensorboardLogger system of using version_{int}.
|
||||
- use_datetime_version (bool): Whether to use a datetime string for version. Defaults to True.
|
||||
- resume_if_exists (bool): Whether this experiment is resuming from a previous run. If True, it sets
|
||||
trainer.resume_from_checkpoint so that the trainer should auto-resume. exp_manager will move files
|
||||
under log_dir to log_dir/run_{int}. Defaults to False. From v1.0.0, when resume_if_exists is True,
|
||||
we would not create version folders to make it easier to find the log folder for next runs.
|
||||
- resume_past_end (bool): exp_manager errors out if resume_if_exists is True and a checkpoint matching
|
||||
*end.ckpt indicating a previous training run fully completed. This behaviour can be disabled, in which
|
||||
case the *end.ckpt will be loaded by setting resume_past_end to True. Defaults to False.
|
||||
- resume_ignore_no_checkpoint (bool): exp_manager errors out if resume_if_exists is True and no checkpoint
|
||||
could be found. This behaviour can be disabled, in which case exp_manager will print a message and
|
||||
continue without restoring, by setting resume_ignore_no_checkpoint to True. Defaults to False.
|
||||
- create_tensorboard_logger (bool): Whether to create a tensorboard logger and attach it to the pytorch
|
||||
lightning trainer. Defaults to True.
|
||||
- summary_writer_kwargs (dict): A dictionary of kwargs that can be passed to lightning's TensorboardLogger
|
||||
class. Note that log_dir is passed by exp_manager and cannot exist in this dict. Defaults to None.
|
||||
- create_wandb_logger (bool): Whether to create a Weights and Baises logger and attach it to the pytorch
|
||||
lightning trainer. Defaults to False.
|
||||
- wandb_logger_kwargs (dict): A dictionary of kwargs that can be passed to lightning's WandBLogger
|
||||
class. Note that name and project are required parameters if create_wandb_logger is True.
|
||||
Defaults to None.
|
||||
- create_checkpoint_callback (bool): Whether to create a ModelCheckpoint callback and attach it to the
|
||||
pytorch lightning trainer. The ModelCheckpoint saves the top 3 models with the best "val_loss", the most
|
||||
recent checkpoint under *last.ckpt, and the final checkpoint after training completes under *end.ckpt.
|
||||
Defaults to True.
|
||||
- files_to_copy (list): A list of files to copy to the experiment logging directory. Defaults to None which
|
||||
copies no files.
|
||||
|
||||
returns:
|
||||
log_dir (Path): The final logging directory where logging files are saved. Usually the concatenation of
|
||||
exp_dir, name, and version.
|
||||
"""
|
||||
# Add rank information to logger
|
||||
# Note: trainer.global_rank and trainer.is_global_zero are not set until trainer.fit, so have to hack around it
|
||||
global_rank = trainer.node_rank * trainer.num_gpus + int(os.environ.get("LOCAL_RANK", 0))
|
||||
logging.rank = global_rank
|
||||
|
||||
if cfg is None:
|
||||
logging.error("exp_manager did not receive a cfg argument. It will be disabled.")
|
||||
return
|
||||
if trainer.fast_dev_run:
|
||||
logging.info("Trainer was called with fast_dev_run. exp_manager will return without any functionality.")
|
||||
return
|
||||
|
||||
# Ensure passed cfg is compliant with ExpManagerConfig
|
||||
schema = OmegaConf.structured(ExpManagerConfig)
|
||||
if isinstance(cfg, dict):
|
||||
cfg = OmegaConf.create(cfg)
|
||||
elif not isinstance(cfg, DictConfig):
|
||||
raise ValueError(f"cfg was type: {type(cfg)}. Expected either a dict or a DictConfig")
|
||||
cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=True))
|
||||
cfg = OmegaConf.merge(schema, cfg)
|
||||
|
||||
error_checks(trainer, cfg) # Ensures that trainer options are compliant with NeMo and exp_manager arguments
|
||||
|
||||
log_dir, exp_dir, name, version = get_log_dir(
|
||||
trainer=trainer,
|
||||
exp_dir=cfg.exp_dir,
|
||||
name=cfg.name,
|
||||
version=cfg.version,
|
||||
explicit_log_dir=cfg.explicit_log_dir,
|
||||
use_datetime_version=cfg.use_datetime_version,
|
||||
resume_if_exists=cfg.resume_if_exists,
|
||||
)
|
||||
|
||||
if cfg.resume_if_exists:
|
||||
check_resume(trainer, log_dir, cfg.resume_past_end, cfg.resume_ignore_no_checkpoint)
|
||||
|
||||
checkpoint_name = name
|
||||
# If name returned from get_log_dir is "", use cfg.name for checkpointing
|
||||
if checkpoint_name is None or checkpoint_name == '':
|
||||
checkpoint_name = cfg.name or "default"
|
||||
cfg.name = name # Used for configure_loggers so that the log_dir is properly set even if name is ""
|
||||
cfg.version = version
|
||||
|
||||
# update app_state with log_dir, exp_dir, etc
|
||||
app_state = AppState()
|
||||
app_state.log_dir = log_dir
|
||||
app_state.exp_dir = exp_dir
|
||||
app_state.name = name
|
||||
app_state.version = version
|
||||
app_state.checkpoint_name = checkpoint_name
|
||||
app_state.create_checkpoint_callback = cfg.create_checkpoint_callback
|
||||
app_state.checkpoint_callback_params = cfg.checkpoint_callback_params
|
||||
|
||||
# Create the logging directory if it does not exist
|
||||
os.makedirs(log_dir, exist_ok=True) # Cannot limit creation to global zero as all ranks write to own log file
|
||||
logging.info(f'Experiments will be logged at {log_dir}')
|
||||
trainer._default_root_dir = log_dir
|
||||
|
||||
# Handle Loggers by creating file and handle DEBUG statements
|
||||
log_file = log_dir / f'nemo_log_globalrank-{global_rank}_localrank-{int(os.environ.get("LOCAL_RANK", 0))}.txt'
|
||||
logging.add_file_handler(log_file)
|
||||
|
||||
# For some reason, LearningRateLogger requires trainer to have a logger. Safer to create logger on all ranks
|
||||
# not just global rank 0.
|
||||
if cfg.create_tensorboard_logger or cfg.create_wandb_logger:
|
||||
configure_loggers(
|
||||
trainer,
|
||||
exp_dir,
|
||||
cfg.name,
|
||||
cfg.version,
|
||||
cfg.create_tensorboard_logger,
|
||||
cfg.summary_writer_kwargs,
|
||||
cfg.create_wandb_logger,
|
||||
cfg.wandb_logger_kwargs,
|
||||
)
|
||||
|
||||
if cfg.create_checkpoint_callback:
|
||||
configure_checkpointing(trainer, log_dir, checkpoint_name, cfg.checkpoint_callback_params)
|
||||
|
||||
if is_global_rank_zero():
|
||||
# Move files_to_copy to folder and add git information if present
|
||||
if cfg.files_to_copy:
|
||||
for _file in cfg.files_to_copy:
|
||||
copy(Path(_file), log_dir)
|
||||
|
||||
# Create files for cmd args and git info
|
||||
with open(log_dir / 'cmd-args.log', 'w') as _file:
|
||||
_file.write(" ".join(sys.argv))
|
||||
|
||||
# Try to get git hash
|
||||
git_repo, git_hash = get_git_hash()
|
||||
if git_repo:
|
||||
with open(log_dir / 'git-info.log', 'w') as _file:
|
||||
_file.write(f'commit hash: {git_hash}')
|
||||
_file.write(get_git_diff())
|
||||
|
||||
# Add err_file logging to global_rank zero
|
||||
logging.add_err_file_handler(log_dir / 'nemo_error_log.txt')
|
||||
|
||||
# Add lightning file logging to global_rank zero
|
||||
add_filehandlers_to_pl_logger(log_dir / 'lightning_logs.txt', log_dir / 'nemo_error_log.txt')
|
||||
|
||||
return log_dir
|
||||
|
||||
|
||||
def error_checks(trainer: 'pytorch_lightning.Trainer', cfg: Optional[Union[DictConfig, Dict]] = None):
|
||||
"""
|
||||
Checks that the passed trainer is compliant with NeMo and exp_manager's passed configuration. Checks that:
|
||||
- Throws error when hydra has changed the working directory. This causes issues with lightning's DDP
|
||||
- Throws error when trainer has loggers defined but create_tensorboard_logger or create_WandB_logger is True
|
||||
- Prints error messages when 1) run on multi-node and not Slurm, and 2) run on multi-gpu without DDP
|
||||
"""
|
||||
if HydraConfig.initialized() and get_original_cwd() != os.getcwd():
|
||||
raise ValueError(
|
||||
"Hydra changed the working directory. This interferes with ExpManger's functionality. Please pass "
|
||||
"hydra.run.dir=. to your python script."
|
||||
)
|
||||
if trainer.logger is not None and (cfg.create_tensorboard_logger or cfg.create_wandb_logger):
|
||||
raise LoggerMisconfigurationError(
|
||||
"The pytorch lightning trainer that was passed to exp_manager contained a logger, and either "
|
||||
f"create_tensorboard_logger: {cfg.create_tensorboard_logger} or create_wandb_logger: "
|
||||
f"{cfg.create_wandb_logger} was set to True. These can only be used if trainer does not already have a"
|
||||
" logger."
|
||||
)
|
||||
if trainer.num_nodes > 1 and not check_slurm(trainer):
|
||||
logging.error(
|
||||
"You are running multi-node training without SLURM handling the processes."
|
||||
" Please note that this is not tested in NeMo and could result in errors."
|
||||
)
|
||||
if trainer.num_gpus > 1 and not trainer.use_ddp:
|
||||
logging.error(
|
||||
"You are running multi-gpu without ddp.Please note that this is not tested in NeMo and could result in "
|
||||
"errors."
|
||||
)
|
||||
|
||||
|
||||
def check_resume(
|
||||
trainer: 'pytorch_lightning.Trainer',
|
||||
log_dir: str,
|
||||
resume_past_end: bool = False,
|
||||
resume_ignore_no_checkpoint: bool = False,
|
||||
):
|
||||
"""Checks that resume=True was used correctly with the arguments pass to exp_manager. Sets
|
||||
trainer.resume_from_checkpoint as necessary.
|
||||
|
||||
Returns:
|
||||
log_dir (Path): the log_dir
|
||||
exp_dir (str): the base exp_dir without name nor version
|
||||
name (str): The name of the experiment
|
||||
version (str): The version of the experiment
|
||||
|
||||
Raises:
|
||||
NotFoundError: If resume is True, resume_ignore_no_checkpoint is False, and checkpoints could not be found.
|
||||
ValueError: If resume is True, and there were more than 1 checkpoint could found.
|
||||
"""
|
||||
if not log_dir:
|
||||
raise ValueError(f"Resuming requires the log_dir {log_dir} to be passed to exp_manager")
|
||||
|
||||
checkpoint_dir = Path(Path(log_dir) / "checkpoints")
|
||||
|
||||
checkpoint = None
|
||||
end_checkpoints = list(checkpoint_dir.rglob("*end.ckpt"))
|
||||
last_checkpoints = list(checkpoint_dir.rglob("*last.ckpt"))
|
||||
if not checkpoint_dir.exists():
|
||||
if resume_ignore_no_checkpoint:
|
||||
logging.warning(
|
||||
f"There was no checkpoint folder at checkpoint_dir :{checkpoint_dir}. Training from scratch."
|
||||
)
|
||||
return
|
||||
else:
|
||||
raise NotFoundError(f"There was no checkpoint folder at checkpoint_dir :{checkpoint_dir}. Cannot resume.")
|
||||
elif len(end_checkpoints) > 0:
|
||||
if resume_past_end:
|
||||
if len(end_checkpoints) > 1:
|
||||
if 'mp_rank' in str(end_checkpoints[0]):
|
||||
checkpoint = end_checkpoints[0]
|
||||
else:
|
||||
raise ValueError(f"Multiple checkpoints {end_checkpoints} that matches *end.ckpt.")
|
||||
logging.info(f"Resuming from {end_checkpoints[0]}")
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Found {end_checkpoints[0]} indicating that the last training run has already completed."
|
||||
)
|
||||
elif not len(last_checkpoints) > 0:
|
||||
if resume_ignore_no_checkpoint:
|
||||
logging.warning(f"There were no checkpoints found in {checkpoint_dir}. Training from scratch.")
|
||||
return
|
||||
else:
|
||||
raise NotFoundError(f"There were no checkpoints found in {checkpoint_dir}. Cannot resume.")
|
||||
elif len(last_checkpoints) > 1:
|
||||
if 'mp_rank' in str(last_checkpoints[0]):
|
||||
checkpoint = last_checkpoints[0]
|
||||
else:
|
||||
raise ValueError(f"Multiple checkpoints {last_checkpoints} that matches *last.ckpt.")
|
||||
else:
|
||||
logging.info(f"Resuming from {last_checkpoints[0]}")
|
||||
checkpoint = last_checkpoints[0]
|
||||
|
||||
trainer.resume_from_checkpoint = str(checkpoint)
|
||||
|
||||
if is_global_rank_zero():
|
||||
# Check to see if any files exist that need to be moved
|
||||
files_to_move = []
|
||||
for child in Path(log_dir).iterdir():
|
||||
if child.is_file():
|
||||
files_to_move.append(child)
|
||||
|
||||
if len(files_to_move) > 0:
|
||||
# Move old files to a new folder
|
||||
other_run_dirs = Path(log_dir).glob("run_*")
|
||||
run_count = 0
|
||||
for fold in other_run_dirs:
|
||||
if fold.is_dir():
|
||||
run_count += 1
|
||||
new_run_dir = Path(Path(log_dir) / f"run_{run_count}")
|
||||
new_run_dir.mkdir()
|
||||
for _file in files_to_move:
|
||||
move(str(_file), str(new_run_dir))
|
||||
|
||||
|
||||
def check_explicit_log_dir(
|
||||
trainer: 'pytorch_lightning.Trainer', explicit_log_dir: [Path, str], exp_dir: str, name: str, version: str
|
||||
) -> (Path, str, str, str):
|
||||
""" Checks that the passed arguments are compatible with explicit_log_dir.
|
||||
|
||||
Returns:
|
||||
log_dir (Path): the log_dir
|
||||
exp_dir (str): the base exp_dir without name nor version
|
||||
name (str): The name of the experiment
|
||||
version (str): The version of the experiment
|
||||
|
||||
Raise:
|
||||
LoggerMisconfigurationError
|
||||
"""
|
||||
if trainer.logger is not None:
|
||||
raise LoggerMisconfigurationError(
|
||||
"The pytorch lightning trainer that was passed to exp_manager contained a logger and explicit_log_dir: "
|
||||
f"{explicit_log_dir} was pass to exp_manager. Please remove the logger from the lightning trainer."
|
||||
)
|
||||
# Checking only (explicit_log_dir) vs (exp_dir and version).
|
||||
# The `name` will be used as the actual name of checkpoint/archive.
|
||||
if exp_dir or version:
|
||||
logging.error(
|
||||
f"exp_manager received explicit_log_dir: {explicit_log_dir} and at least one of exp_dir: {exp_dir}, "
|
||||
f"or version: {version}. Please note that exp_dir, name, and version will be ignored."
|
||||
)
|
||||
if is_global_rank_zero() and Path(explicit_log_dir).exists():
|
||||
logging.warning(f"Exp_manager is logging to {explicit_log_dir}, but it already exists.")
|
||||
return Path(explicit_log_dir), str(explicit_log_dir), "", ""
|
||||
|
||||
|
||||
def get_log_dir(
|
||||
trainer: 'pytorch_lightning.Trainer',
|
||||
exp_dir: str = None,
|
||||
name: str = None,
|
||||
version: str = None,
|
||||
explicit_log_dir: str = None,
|
||||
use_datetime_version: bool = True,
|
||||
resume_if_exists: bool = False,
|
||||
) -> (Path, str, str, str):
|
||||
"""
|
||||
Obtains the log_dir used for exp_manager.
|
||||
|
||||
Returns:
|
||||
log_dir (Path): the log_dir
|
||||
exp_dir (str): the base exp_dir without name nor version
|
||||
name (str): The name of the experiment
|
||||
version (str): The version of the experiment
|
||||
explicit_log_dir (str): The explicit path to the log folder. Defaults to False.
|
||||
use_datetime_version (bool): Uses date and time as the version of the log folder. Defaults to True.
|
||||
resume_if_exists (bool): if resume_if_exists of the exp_manager's config is enabled or not. When enabled, the
|
||||
version folders would not get created.
|
||||
|
||||
Raise:
|
||||
LoggerMisconfigurationError: If trainer is incompatible with arguments
|
||||
NotFoundError: If resume is True, resume_ignore_no_checkpoint is False, and checkpoints could not be found.
|
||||
ValueError: If resume is True, and there were more than 1 checkpoint could found.
|
||||
"""
|
||||
if explicit_log_dir: # If explicit log_dir was passed, short circuit
|
||||
return check_explicit_log_dir(trainer, explicit_log_dir, exp_dir, name, version)
|
||||
|
||||
# Default exp_dir to ./nemo_experiments if None was passed
|
||||
_exp_dir = exp_dir
|
||||
if exp_dir is None:
|
||||
_exp_dir = str(Path.cwd() / 'nemo_experiments')
|
||||
|
||||
# If the user has already defined a logger for the trainer, use the logger defaults for logging directory
|
||||
if trainer.logger is not None:
|
||||
if trainer.logger.save_dir:
|
||||
if exp_dir:
|
||||
raise LoggerMisconfigurationError(
|
||||
"The pytorch lightning trainer that was passed to exp_manager contained a logger, the logger's "
|
||||
f"save_dir was not None, and exp_dir ({exp_dir}) was not None. If trainer.logger.save_dir "
|
||||
"exists, exp_manager will use trainer.logger.save_dir as the logging directory and exp_dir "
|
||||
"must be None."
|
||||
)
|
||||
_exp_dir = trainer.logger.save_dir
|
||||
if name:
|
||||
raise LoggerMisconfigurationError(
|
||||
"The pytorch lightning trainer that was passed to exp_manager contained a logger, and name: "
|
||||
f"{name} was also passed to exp_manager. If the trainer contains a "
|
||||
"logger, exp_manager will use trainer.logger.name, and name passed to exp_manager must be None."
|
||||
)
|
||||
name = trainer.logger.name
|
||||
version = f"version_{trainer.logger.version}"
|
||||
# Use user-defined exp_dir, project_name, exp_name, and versioning options
|
||||
else:
|
||||
name = name or "default"
|
||||
version = version or os.environ.get(NEMO_ENV_VARNAME_VERSION, None)
|
||||
|
||||
if not version:
|
||||
if resume_if_exists:
|
||||
logging.warning(
|
||||
"No version folders would be created under the log folder as 'resume_if_exists' is enabled."
|
||||
)
|
||||
version = None
|
||||
elif is_global_rank_zero():
|
||||
if use_datetime_version:
|
||||
version = time.strftime('%Y-%m-%d_%H-%M-%S')
|
||||
else:
|
||||
tensorboard_logger = TensorBoardLogger(save_dir=Path(_exp_dir), name=name, version=version)
|
||||
version = f"version_{tensorboard_logger.version}"
|
||||
os.environ[NEMO_ENV_VARNAME_VERSION] = "" if version is None else version
|
||||
|
||||
log_dir = Path(_exp_dir) / Path(str(name)) / Path("" if version is None else str(version))
|
||||
return log_dir, str(_exp_dir), name, version
|
||||
|
||||
|
||||
def get_git_hash():
|
||||
"""
|
||||
Helper function that tries to get the commit hash if running inside a git folder
|
||||
|
||||
returns:
|
||||
Bool: Whether the git subprocess ran without error
|
||||
str: git subprocess output or error message
|
||||
"""
|
||||
try:
|
||||
return (
|
||||
True,
|
||||
subprocess.check_output(['git', 'rev-parse', 'HEAD'], stderr=subprocess.STDOUT).decode(),
|
||||
)
|
||||
except subprocess.CalledProcessError as err:
|
||||
return False, "{}\n".format(err.output.decode("utf-8"))
|
||||
|
||||
|
||||
def get_git_diff():
|
||||
"""
|
||||
Helper function that tries to get the git diff if running inside a git folder
|
||||
|
||||
returns:
|
||||
Bool: Whether the git subprocess ran without error
|
||||
str: git subprocess output or error message
|
||||
"""
|
||||
try:
|
||||
return subprocess.check_output(['git', 'diff'], stderr=subprocess.STDOUT).decode()
|
||||
except subprocess.CalledProcessError as err:
|
||||
return "{}\n".format(err.output.decode("utf-8"))
|
||||
|
||||
|
||||
class LoggerList(_LoggerCollection):
|
||||
""" A thin wrapper on Lightning's LoggerCollection such that name and version are better aligned with exp_manager
|
||||
"""
|
||||
|
||||
def __init__(self, _logger_iterable, nemo_name=None, nemo_version=""):
|
||||
super().__init__(_logger_iterable)
|
||||
self._nemo_name = nemo_name
|
||||
self._nemo_version = nemo_version
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._nemo_name
|
||||
|
||||
@property
|
||||
def version(self) -> str:
|
||||
return self._nemo_version
|
||||
|
||||
|
||||
def configure_loggers(
|
||||
trainer: 'pytorch_lightning.Trainer',
|
||||
exp_dir: [Path, str],
|
||||
name: str,
|
||||
version: str,
|
||||
create_tensorboard_logger: bool,
|
||||
summary_writer_kwargs: dict,
|
||||
create_wandb_logger: bool,
|
||||
wandb_kwargs: dict,
|
||||
):
|
||||
""" Creates TensorboardLogger and/or WandBLogger and attach them to trainer. Raises ValueError if
|
||||
summary_writer_kwargs or wandb_kwargs are misconfigured.
|
||||
"""
|
||||
# Potentially create tensorboard logger and/or WandBLogger
|
||||
logger_list = []
|
||||
if create_tensorboard_logger:
|
||||
if summary_writer_kwargs is None:
|
||||
summary_writer_kwargs = {}
|
||||
elif "log_dir" in summary_writer_kwargs:
|
||||
raise ValueError(
|
||||
"You cannot pass `log_dir` as part of `summary_writer_kwargs`. `log_dir` is handled by lightning's "
|
||||
"TensorBoardLogger logger."
|
||||
)
|
||||
tensorboard_logger = TensorBoardLogger(save_dir=exp_dir, name=name, version=version, **summary_writer_kwargs)
|
||||
logger_list.append(tensorboard_logger)
|
||||
logging.info("TensorboardLogger has been set up")
|
||||
|
||||
if create_wandb_logger:
|
||||
if wandb_kwargs is None:
|
||||
wandb_kwargs = {}
|
||||
if "name" not in wandb_kwargs and "project" not in wandb_kwargs:
|
||||
raise ValueError("name and project are required for wandb_logger")
|
||||
wandb_logger = WandbLogger(save_dir=exp_dir, version=version, **wandb_kwargs)
|
||||
|
||||
logger_list.append(wandb_logger)
|
||||
logging.info("WandBLogger has been set up")
|
||||
|
||||
logger_list = (
|
||||
LoggerList(logger_list, nemo_name=name, nemo_version=version) if len(logger_list) > 1 else logger_list[0]
|
||||
)
|
||||
trainer.logger_connector.configure_logger(logger_list)
|
||||
|
||||
|
||||
class NeMoModelCheckpoint(ModelCheckpoint):
|
||||
""" Light wrapper around Lightning's ModelCheckpoint to force a saved checkpoint on train_end
|
||||
"""
|
||||
|
||||
def __init__(self, always_save_nemo=False, save_best_model=False, postfix=".nemo", **kwargs):
|
||||
# Parse and store "extended" parameters: save_best model and postfix.
|
||||
self.always_save_nemo = always_save_nemo
|
||||
self.save_best_model = save_best_model
|
||||
self.postfix = postfix
|
||||
self.previous_best_path = ""
|
||||
|
||||
# `prefix` is deprecated
|
||||
if 'prefix' in kwargs:
|
||||
self.prefix = kwargs.pop('prefix')
|
||||
else:
|
||||
self.prefix = ""
|
||||
|
||||
# Call the parent class constructor with the remaining kwargs.
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@rank_zero_only
|
||||
def on_save_checkpoint(self, trainer, pl_module, checkpoint):
|
||||
output = super().on_save_checkpoint(trainer, pl_module, checkpoint)
|
||||
|
||||
if not self.always_save_nemo:
|
||||
return output
|
||||
|
||||
# Load the best model and then re-save it
|
||||
app_state = AppState()
|
||||
# since we are creating tarfile artifacts we need to update .nemo path
|
||||
app_state.model_restore_path = os.path.abspath(
|
||||
os.path.expanduser(os.path.join(self.dirpath, self.prefix + self.postfix))
|
||||
)
|
||||
if self.save_best_model:
|
||||
if not os.path.exists(self.best_model_path):
|
||||
return output
|
||||
|
||||
if self.best_model_path == self.previous_best_path:
|
||||
return output
|
||||
|
||||
self.previous_model_path = self.best_model_path
|
||||
old_state_dict = deepcopy(pl_module.state_dict())
|
||||
checkpoint = torch.load(self.best_model_path, map_location='cpu')
|
||||
if 'state_dict' in checkpoint:
|
||||
checkpoint = checkpoint['state_dict']
|
||||
# get a new instanace of the model
|
||||
pl_module.load_state_dict(checkpoint, strict=True)
|
||||
pl_module.save_to(save_path=app_state.model_restore_path)
|
||||
pl_module.load_state_dict(old_state_dict, strict=True)
|
||||
else:
|
||||
pl_module.save_to(save_path=app_state.model_restore_path)
|
||||
return output
|
||||
|
||||
@rank_zero_only
|
||||
def on_train_end(self, trainer, pl_module):
|
||||
if trainer.fast_dev_run:
|
||||
return None
|
||||
app_state = AppState()
|
||||
if app_state.model_parallel_size is not None:
|
||||
return None
|
||||
|
||||
# TODO: make this work for model parallel, need to call on data parallel rank 0 and update best_model_path
|
||||
# Load the best model and then re-save it
|
||||
if self.save_best_model:
|
||||
trainer.checkpoint_connector.restore(self.best_model_path, on_gpu=trainer.on_gpu)
|
||||
pl_module.save_to(save_path=os.path.join(self.dirpath, self.prefix + self.postfix))
|
||||
|
||||
|
||||
def configure_checkpointing(trainer: 'pytorch_lightning.Trainer', log_dir: Path, name: str, params: 'DictConfig'):
|
||||
""" Adds ModelCheckpoint to trainer. Raises CheckpointMisconfigurationError if trainer already has a ModelCheckpoint
|
||||
callback or if trainer.weights_save_path was passed to Trainer.
|
||||
"""
|
||||
for callback in trainer.callbacks:
|
||||
if isinstance(callback, ModelCheckpoint):
|
||||
raise CheckpointMisconfigurationError(
|
||||
"The pytorch lightning trainer that was passed to exp_manager contained a ModelCheckpoint "
|
||||
"and create_checkpoint_callback was set to True. Please either set create_checkpoint_callback "
|
||||
"to False, or remove ModelCheckpoint from the lightning trainer"
|
||||
)
|
||||
if Path(trainer.weights_save_path) != Path.cwd():
|
||||
raise CheckpointMisconfigurationError(
|
||||
"The pytorch lightning was passed weights_save_path. This variable is ignored by exp_manager"
|
||||
)
|
||||
|
||||
# Create the callback and attach it to trainer
|
||||
if "filepath" in params:
|
||||
if params.filepath is not None:
|
||||
logging.warning("filepath is deprecated. Please switch to dirpath and filename instead")
|
||||
if params.dirpath is None:
|
||||
params.dirpath = Path(params.filepath).parent
|
||||
if params.filename is None:
|
||||
params.filename = Path(params.filepath).name
|
||||
with open_dict(params):
|
||||
del params["filepath"]
|
||||
if params.dirpath is None:
|
||||
params.dirpath = Path(log_dir / 'checkpoints')
|
||||
if params.filename is None:
|
||||
params.filename = f'{name}--{{{params.monitor}:.2f}}-{{epoch}}'
|
||||
if params.prefix is None:
|
||||
params.prefix = name
|
||||
NeMoModelCheckpoint.CHECKPOINT_NAME_LAST = params.filename + '-last'
|
||||
|
||||
logging.debug(params.dirpath)
|
||||
logging.debug(params.filename)
|
||||
logging.debug(params.prefix)
|
||||
|
||||
if "val" in params.monitor:
|
||||
if (
|
||||
trainer.max_epochs is not None
|
||||
and trainer.max_epochs != -1
|
||||
and trainer.max_epochs < trainer.check_val_every_n_epoch
|
||||
):
|
||||
logging.error(
|
||||
"The checkpoint callback was told to monitor a validation value but trainer.max_epochs("
|
||||
f"{trainer.max_epochs}) was less than trainer.check_val_every_n_epoch({trainer.check_val_every_n_epoch}"
|
||||
f"). It is very likely this run will fail with ModelCheckpoint(monitor='{params.monitor}') not found "
|
||||
"in the returned metrics. Please ensure that validation is run within trainer.max_epochs."
|
||||
)
|
||||
elif trainer.max_steps is not None:
|
||||
logging.warning(
|
||||
"The checkpoint callback was told to monitor a validation value and trainer's max_steps was set to "
|
||||
f"{trainer.max_steps}. Please ensure that max_steps will run for at least "
|
||||
f"{trainer.check_val_every_n_epoch} epochs to ensure that checkpointing will not error out."
|
||||
)
|
||||
|
||||
checkpoint_callback = NeMoModelCheckpoint(**params)
|
||||
checkpoint_callback.last_model_path = trainer.resume_from_checkpoint or ""
|
||||
trainer.callbacks.append(checkpoint_callback)
|
||||
|
||||
|
||||
def check_slurm(trainer):
|
||||
try:
|
||||
return trainer.accelerator_connector.is_slurm_managing_tasks
|
||||
except AttributeError:
|
||||
return False
|
||||
@@ -0,0 +1,225 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Callable, Dict, Optional, Type
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
apex_available = True
|
||||
|
||||
try:
|
||||
from apex.normalization.fused_layer_norm import FusedLayerNorm
|
||||
|
||||
def replace_FusedLayerNorm(n: nn.Module) -> Optional[nn.BatchNorm2d]:
|
||||
"""
|
||||
Replaces Apex's FusedLayerNorm with nn.LayerNorm. This is required for ONNX export.
|
||||
Args:
|
||||
n: the FusedLayerNorm pytorch module to replace
|
||||
Returns:
|
||||
Equivalent LayerNorm module
|
||||
"""
|
||||
if not apex_available or not isinstance(n, FusedLayerNorm):
|
||||
return None
|
||||
|
||||
dev = next(n.parameters()).device
|
||||
mod = nn.LayerNorm(n.normalized_shape, eps=n.eps, elementwise_affine=n.elementwise_affine,).to(dev)
|
||||
|
||||
n_state = n.state_dict()
|
||||
mod.load_state_dict(n_state)
|
||||
return mod
|
||||
|
||||
default_Apex_replacements = {"FusedLayerNorm": replace_FusedLayerNorm}
|
||||
|
||||
except Exception as e:
|
||||
default_Apex_replacements = {}
|
||||
apex_available = False
|
||||
|
||||
|
||||
def expand_Conv1D(conv1d: nn.Module) -> Optional[nn.Conv2d]:
|
||||
"""
|
||||
Expands a Conv1D into a Conv2D. This is required for many (closed source) commercial tools with poor support for 1D Convolutions in Onnx.
|
||||
Args:
|
||||
conv1d: the Conv1D pytorch module to expand
|
||||
Returns:
|
||||
conv2d: Conv2D module with identical weights and params
|
||||
"""
|
||||
if not isinstance(conv1d, nn.Conv1d):
|
||||
return None
|
||||
conv2d = nn.Conv2d(
|
||||
conv1d.in_channels,
|
||||
conv1d.out_channels,
|
||||
kernel_size=(conv1d.kernel_size[0], 1),
|
||||
stride=(conv1d.stride[0], 1),
|
||||
padding=(conv1d.padding[0], 0),
|
||||
dilation=(conv1d.dilation[0], 1),
|
||||
groups=conv1d.groups,
|
||||
padding_mode=conv1d.padding_mode,
|
||||
).to(device=conv1d.weight.device, dtype=conv1d.weight.dtype)
|
||||
conv2d.bias = conv1d.bias
|
||||
conv2d.weight = nn.Parameter(conv1d.weight.unsqueeze(-1))
|
||||
# check that expansion is valid
|
||||
for _ in range(2):
|
||||
sample_input = torch.rand(1, conv1d.in_channels, 256).to(
|
||||
device=conv1d.weight.device, dtype=conv1d.weight.dtype
|
||||
)
|
||||
close = conv1d(sample_input).mean() - conv2d(sample_input.unsqueeze(-1)).squeeze().mean()
|
||||
if close.abs() > 1.0:
|
||||
raise ValueError("Unable to expand Conv1D to Conv2D")
|
||||
return conv2d
|
||||
|
||||
|
||||
def expand_BatchNorm1d(bn1d: nn.Module) -> Optional[nn.BatchNorm2d]:
|
||||
"""
|
||||
Expands a BatchNorm1d into a BatchNorm2d. This is required for many (closed source) commercial tools with poor support for BatchNorm1d in Onnx.
|
||||
Args:
|
||||
bn1d: the BatchNorm1d pytorch module to expand
|
||||
Returns:
|
||||
bn2d: BatchNorm2d module with identical weights and params
|
||||
"""
|
||||
if not isinstance(bn1d, nn.BatchNorm1d):
|
||||
return None
|
||||
mod = torch.nn.BatchNorm2d(
|
||||
bn1d.num_features,
|
||||
eps=bn1d.eps,
|
||||
momentum=bn1d.momentum,
|
||||
affine=bn1d.affine,
|
||||
track_running_stats=bn1d.track_running_stats,
|
||||
).to(device=conv1d.weight.device, dtype=conv1d.weight.dtype)
|
||||
bn_state = bn1d.state_dict()
|
||||
mod.load_state_dict(bn_state)
|
||||
return mod
|
||||
|
||||
|
||||
def expand_ConvTranspose1D(conv1d: nn.Module) -> Optional[nn.ConvTranspose2d]:
|
||||
"""
|
||||
Expands a Conv1D into a Conv2D. This is required for many (closed source) commercial tools with poor support for 1D Convolutions in Onnx.
|
||||
Args:
|
||||
conv1d: the Conv1D pytorch module to expand
|
||||
Returns:
|
||||
conv2d: Conv2D module with identical weights and params
|
||||
"""
|
||||
if not isinstance(conv1d, nn.ConvTranspose1d):
|
||||
return None
|
||||
conv2d = nn.ConvTranspose2d(
|
||||
conv1d.in_channels,
|
||||
conv1d.out_channels,
|
||||
kernel_size=(conv1d.kernel_size[0], 1),
|
||||
stride=(conv1d.stride[0], 1),
|
||||
padding=(conv1d.padding[0], 0),
|
||||
dilation=(conv1d.dilation[0], 1),
|
||||
groups=conv1d.groups,
|
||||
padding_mode=conv1d.padding_mode,
|
||||
).to(device=conv1d.weight.device, dtype=conv1d.weight.dtype)
|
||||
conv2d.bias = conv1d.bias
|
||||
conv2d.weight = nn.Parameter(conv1d.weight.unsqueeze(-1))
|
||||
# check that expansion is valid
|
||||
for _ in range(2):
|
||||
sample_input = torch.rand(1, conv1d.in_channels, 256).to(
|
||||
device=conv1d.weight.device, dtype=conv1d.weight.dtype
|
||||
)
|
||||
close = conv1d(sample_input).mean() - conv2d(sample_input.unsqueeze(-1)).squeeze().mean()
|
||||
if close.abs() > 1.0:
|
||||
raise ValueError("Unable to expand Conv1D to Conv2D")
|
||||
return conv2d
|
||||
|
||||
|
||||
def simple_replace(BaseT: Type[nn.Module], DestT: Type[nn.Module]) -> Callable[[nn.Module], Optional[nn.Module]]:
|
||||
"""
|
||||
Generic function generator to replace BaseT module with DestT. BaseT and DestT should have same atrributes. No weights are copied.
|
||||
Args:
|
||||
BaseT : module type to replace
|
||||
DestT : destination module type
|
||||
Returns:
|
||||
swap function to replace BaseT module with DestT
|
||||
"""
|
||||
|
||||
def expansion_fn(mod: nn.Module) -> Optional[nn.Module]:
|
||||
if not isinstance(mod, BaseT):
|
||||
return None
|
||||
args = [getattr(mod, name, None) for name in mod.__constants__]
|
||||
out = DestT(*args)
|
||||
return out
|
||||
|
||||
return expansion_fn
|
||||
|
||||
|
||||
def swap_modules(model: nn.Module, mapping: Dict[str, nn.Module]):
|
||||
"""
|
||||
This function swaps nested modules as specified by "dot paths" in mod with a desired replacement. This allows
|
||||
for swapping nested modules through arbitrary levels if children
|
||||
|
||||
NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
|
||||
|
||||
"""
|
||||
for path, new_mod in mapping.items():
|
||||
expanded_path = path.split(".")
|
||||
parent_mod = model
|
||||
for sub_path in expanded_path[:-1]:
|
||||
parent_mod = parent_mod._modules[sub_path] # noqa
|
||||
parent_mod._modules[expanded_path[-1]] = new_mod # noqa
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def replace_modules(
|
||||
model: nn.Module, expansions: Dict[str, Callable[[nn.Module], Optional[nn.Module]]] = None
|
||||
) -> nn.Module:
|
||||
"""
|
||||
Top-level function to replace modules in model, specified by class name with a desired replacement.
|
||||
NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
|
||||
Args:
|
||||
model : top level module
|
||||
expansions : replacement dictionary: module class name -> replacement function generator
|
||||
Returns:
|
||||
model, possibly modified in-place
|
||||
"""
|
||||
mapping: Dict[str, nn.Module] = {}
|
||||
for name, m in model.named_modules():
|
||||
m_type = type(m).__name__
|
||||
if m_type in expansions:
|
||||
swapped = expansions[m_type](m)
|
||||
if swapped:
|
||||
mapping[name] = swapped
|
||||
logging.warning(f"Swapped {len(mapping)} modules")
|
||||
swap_modules(model, mapping)
|
||||
return model
|
||||
|
||||
|
||||
default_1D_2D_replacements = {
|
||||
"Conv1d": expand_Conv1D,
|
||||
"ConvTranspose1d": expand_ConvTranspose1D,
|
||||
"BatchNorm1d": expand_BatchNorm1d,
|
||||
"AdaptiveAvgPool1d": simple_replace(nn.AdaptiveAvgPool1d, nn.AdaptiveAvgPool2d),
|
||||
"AvgPool1d": simple_replace(nn.AvgPool1d, nn.AvgPool2d),
|
||||
}
|
||||
|
||||
|
||||
def replace_for_export(model: nn.Module, replace_1D_2D: bool = False) -> nn.Module:
|
||||
"""
|
||||
Top-level function to replace default set of modules in model
|
||||
NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
|
||||
Args:
|
||||
model : top level module
|
||||
replace_1D_2D : include 1D -> 2D replacements
|
||||
Returns:
|
||||
model, possibly modified in-place
|
||||
"""
|
||||
replace_modules(model, default_Apex_replacements)
|
||||
if replace_1D_2D:
|
||||
# TODO: add squeeze/unsqueeze
|
||||
replace_modules(model, default_1D_2D_replacements)
|
||||
@@ -0,0 +1,13 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
@@ -0,0 +1,135 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import logging
|
||||
|
||||
from nemo.utils.formatters.colors import Fore as ForegroundColors
|
||||
from nemo.utils.formatters.utils import check_color_support, to_unicode
|
||||
|
||||
__all__ = ["BaseNeMoFormatter"]
|
||||
|
||||
|
||||
class BaseFormatter(logging.Formatter):
|
||||
"""
|
||||
Log formatter used in Tornado. Key features of this formatter are:
|
||||
* Color support when logging to a terminal that supports it.
|
||||
* Timestamps on every log line.
|
||||
* Robust against str/bytes encoding problems.
|
||||
"""
|
||||
|
||||
DEFAULT_FORMAT = "%(color)s[%(levelname)1.1s %(asctime)s %(module)s:%(lineno)d]%(end_color)s %(message)s"
|
||||
|
||||
DEFAULT_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
|
||||
|
||||
DEFAULT_COLORS = {
|
||||
logging.DEBUG: ForegroundColors.CYAN,
|
||||
logging.INFO: ForegroundColors.GREEN,
|
||||
logging.WARNING: ForegroundColors.YELLOW,
|
||||
logging.ERROR: ForegroundColors.MAGENTA,
|
||||
logging.CRITICAL: ForegroundColors.RED,
|
||||
}
|
||||
|
||||
def __init__(self, color=True, fmt=None, datefmt=None, colors=None):
|
||||
r"""
|
||||
:arg bool color: Enables color support.
|
||||
:arg string fmt: Log message format.
|
||||
It will be applied to the attributes dict of log records. The
|
||||
text between ``%(color)s`` and ``%(end_color)s`` will be colored
|
||||
depending on the level if color support is on.
|
||||
:arg dict colors: color mappings from logging level to terminal color
|
||||
code
|
||||
:arg string datefmt: Datetime format.
|
||||
Used for formatting ``(asctime)`` placeholder in ``prefix_fmt``.
|
||||
.. versionchanged:: 3.2
|
||||
Added ``fmt`` and ``datefmt`` arguments.
|
||||
"""
|
||||
|
||||
if fmt is None:
|
||||
fmt = self.DEFAULT_FORMAT
|
||||
|
||||
if datefmt is None:
|
||||
datefmt = self.DEFAULT_DATE_FORMAT
|
||||
|
||||
if colors is None:
|
||||
colors = self.DEFAULT_COLORS
|
||||
|
||||
logging.Formatter.__init__(self, datefmt=datefmt)
|
||||
|
||||
self._fmt = fmt
|
||||
self._colors = {}
|
||||
self._normal = ""
|
||||
|
||||
if color and check_color_support():
|
||||
self._colors = colors
|
||||
self._normal = ForegroundColors.RESET
|
||||
|
||||
def format(self, record):
|
||||
try:
|
||||
message = record.getMessage()
|
||||
assert isinstance(message, str) # guaranteed by logging
|
||||
# Encoding notes: The logging module prefers to work with character
|
||||
# strings, but only enforces that log messages are instances of
|
||||
# basestring. In python 2, non-ascii bytestrings will make
|
||||
# their way through the logging framework until they blow up with
|
||||
# an unhelpful decoding error (with this formatter it happens
|
||||
# when we attach the prefix, but there are other opportunities for
|
||||
# exceptions further along in the framework).
|
||||
#
|
||||
# If a byte string makes it this far, convert it to unicode to
|
||||
# ensure it will make it out to the logs. Use repr() as a fallback
|
||||
# to ensure that all byte strings can be converted successfully,
|
||||
# but don't do it by default so we don't add extra quotes to ascii
|
||||
# bytestrings. This is a bit of a hacky place to do this, but
|
||||
# it's worth it since the encoding errors that would otherwise
|
||||
# result are so useless (and tornado is fond of using utf8-encoded
|
||||
# byte strings wherever possible).
|
||||
record.message = to_unicode(message)
|
||||
|
||||
except Exception as e:
|
||||
record.message = "Bad message (%r): %r" % (e, record.__dict__)
|
||||
|
||||
record.asctime = self.formatTime(record, self.datefmt)
|
||||
|
||||
if record.levelno in self._colors:
|
||||
record.color = self._colors[record.levelno]
|
||||
record.end_color = self._normal
|
||||
else:
|
||||
record.color = record.end_color = ""
|
||||
|
||||
formatted = self._fmt % record.__dict__
|
||||
|
||||
if record.exc_info:
|
||||
if not record.exc_text:
|
||||
record.exc_text = self.formatException(record.exc_info)
|
||||
|
||||
if record.exc_text:
|
||||
# exc_text contains multiple lines. We need to _safe_unicode
|
||||
# each line separately so that non-utf8 bytes don't cause
|
||||
# all the newlines to turn into '\n'.
|
||||
lines = [formatted.rstrip()]
|
||||
lines.extend(to_unicode(ln) for ln in record.exc_text.split("\n"))
|
||||
|
||||
formatted = "\n".join(lines)
|
||||
return formatted.replace("\n", "\n ")
|
||||
|
||||
|
||||
class BaseNeMoFormatter(BaseFormatter):
|
||||
DEFAULT_FORMAT = "%(color)s[NeMo %(levelname)1.1s %(asctime)s %(module)s:%(lineno)d]%(end_color)s %(message)s"
|
||||
|
||||
|
||||
class DebugNeMoFormatter(BaseFormatter):
|
||||
DEFAULT_FORMAT = (
|
||||
"%(color)s[NeMo %(levelname)1.1s %(asctime)s %(module)s:%(lineno)d rank:%(rank)s]%(end_color)s %(message)s"
|
||||
)
|
||||
@@ -0,0 +1,121 @@
|
||||
# Source: https://github.com/tartley/colorama/blob/master/colorama/ansi.py
|
||||
# Copyright: Jonathan Hartley 2013. BSD 3-Clause license.
|
||||
#
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
CSI = "\033["
|
||||
OSC = "\033]"
|
||||
BEL = "\007"
|
||||
|
||||
|
||||
def code_to_chars(code):
|
||||
return CSI + str(code) + "m"
|
||||
|
||||
|
||||
def set_title(title):
|
||||
return OSC + "2;" + title + BEL
|
||||
|
||||
|
||||
def clear_screen(mode=2):
|
||||
return CSI + str(mode) + "J"
|
||||
|
||||
|
||||
def clear_line(mode=2):
|
||||
return CSI + str(mode) + "K"
|
||||
|
||||
|
||||
class AnsiCodes(object):
|
||||
def __init__(self):
|
||||
# the subclasses declare class attributes which are numbers.
|
||||
# Upon instantiation we define instance attributes, which are the same
|
||||
# as the class attributes but wrapped with the ANSI escape sequence
|
||||
for name in dir(self):
|
||||
if not name.startswith("_"):
|
||||
value = getattr(self, name)
|
||||
setattr(self, name, code_to_chars(value))
|
||||
|
||||
|
||||
class AnsiCursor(object):
|
||||
def UP(self, n=1):
|
||||
return CSI + str(n) + "A"
|
||||
|
||||
def DOWN(self, n=1):
|
||||
return CSI + str(n) + "B"
|
||||
|
||||
def FORWARD(self, n=1):
|
||||
return CSI + str(n) + "C"
|
||||
|
||||
def BACK(self, n=1):
|
||||
return CSI + str(n) + "D"
|
||||
|
||||
def POS(self, x=1, y=1):
|
||||
return CSI + str(y) + ";" + str(x) + "H"
|
||||
|
||||
|
||||
class AnsiFore(AnsiCodes):
|
||||
BLACK = 30
|
||||
RED = 31
|
||||
GREEN = 32
|
||||
YELLOW = 33
|
||||
BLUE = 34
|
||||
MAGENTA = 35
|
||||
CYAN = 36
|
||||
WHITE = 37
|
||||
RESET = 39
|
||||
|
||||
# These are fairly well supported, but not part of the standard.
|
||||
LIGHTBLACK_EX = 90
|
||||
LIGHTRED_EX = 91
|
||||
LIGHTGREEN_EX = 92
|
||||
LIGHTYELLOW_EX = 93
|
||||
LIGHTBLUE_EX = 94
|
||||
LIGHTMAGENTA_EX = 95
|
||||
LIGHTCYAN_EX = 96
|
||||
LIGHTWHITE_EX = 97
|
||||
|
||||
|
||||
class AnsiBack(AnsiCodes):
|
||||
BLACK = 40
|
||||
RED = 41
|
||||
GREEN = 42
|
||||
YELLOW = 43
|
||||
BLUE = 44
|
||||
MAGENTA = 45
|
||||
CYAN = 46
|
||||
WHITE = 47
|
||||
RESET = 49
|
||||
|
||||
# These are fairly well supported, but not part of the standard.
|
||||
LIGHTBLACK_EX = 100
|
||||
LIGHTRED_EX = 101
|
||||
LIGHTGREEN_EX = 102
|
||||
LIGHTYELLOW_EX = 103
|
||||
LIGHTBLUE_EX = 104
|
||||
LIGHTMAGENTA_EX = 105
|
||||
LIGHTCYAN_EX = 106
|
||||
LIGHTWHITE_EX = 107
|
||||
|
||||
|
||||
class AnsiStyle(AnsiCodes):
|
||||
BRIGHT = 1
|
||||
DIM = 2
|
||||
NORMAL = 22
|
||||
RESET_ALL = 0
|
||||
|
||||
|
||||
Fore = AnsiFore()
|
||||
Back = AnsiBack()
|
||||
Style = AnsiStyle()
|
||||
Cursor = AnsiCursor()
|
||||
@@ -0,0 +1,46 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import sys
|
||||
|
||||
from nemo.constants import NEMO_ENV_VARNAME_ENABLE_COLORING
|
||||
from nemo.utils.env_var_parsing import get_envbool
|
||||
|
||||
__all__ = ["check_color_support", "to_unicode"]
|
||||
|
||||
|
||||
def check_color_support():
|
||||
# Colors can be forced with an env variable
|
||||
if not sys.platform.lower().startswith("win") and get_envbool(NEMO_ENV_VARNAME_ENABLE_COLORING, False):
|
||||
return True
|
||||
|
||||
|
||||
def to_unicode(value):
|
||||
"""
|
||||
Converts a string argument to a unicode string.
|
||||
If the argument is already a unicode string or None, it is returned
|
||||
unchanged. Otherwise it must be a byte string and is decoded as utf8.
|
||||
"""
|
||||
try:
|
||||
if isinstance(value, (str, type(None))):
|
||||
return value
|
||||
|
||||
if not isinstance(value, bytes):
|
||||
raise TypeError("Expected bytes, unicode, or None; got %r" % type(value))
|
||||
|
||||
return value.decode("utf-8")
|
||||
|
||||
except UnicodeDecodeError:
|
||||
return repr(value)
|
||||
@@ -0,0 +1,48 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from nemo.utils.env_var_parsing import get_envint
|
||||
|
||||
|
||||
def is_global_rank_zero():
|
||||
""" Helper function to determine if the current process is global_rank 0 (the main process)
|
||||
"""
|
||||
# Try to get the pytorch RANK env var
|
||||
# RANK is set by torch.distributed.launch
|
||||
rank = get_envint("RANK", None)
|
||||
if rank:
|
||||
return rank == 0
|
||||
|
||||
# If not set by pytorch, we need to determine node_rank
|
||||
def get_node_rank():
|
||||
# Use an equivalent of pytorch lightning's determine_ddp_node_rank()
|
||||
node_rank = 0
|
||||
# First check if running on a slurm cluster
|
||||
# TODO: This check could probably be better
|
||||
num_slurm_tasks = get_envint("SLURM_NTASKS", 0)
|
||||
if num_slurm_tasks > 0:
|
||||
node_rank = get_envint("SLURM_NODEID", 0)
|
||||
else:
|
||||
node_rank_env = get_envint("NODE_RANK", None)
|
||||
group_rank = get_envint("GROUP_RANK", None)
|
||||
if group_rank:
|
||||
node_rank = group_rank
|
||||
# Take from NODE_RANK whenever available
|
||||
if node_rank_env:
|
||||
node_rank = node_rank_env
|
||||
return node_rank
|
||||
|
||||
node_rank = get_node_rank()
|
||||
local_rank = get_envint("LOCAL_RANK", 0)
|
||||
return node_rank == 0 and local_rank == 0
|
||||
@@ -0,0 +1,58 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging as _logging
|
||||
from logging.handlers import MemoryHandler
|
||||
|
||||
import pytorch_lightning as pl
|
||||
|
||||
HANDLERS = {}
|
||||
PATCHED = False
|
||||
|
||||
|
||||
def add_memory_handlers_to_pl_logger():
|
||||
"""
|
||||
Adds two MemoryHandlers to pytorch_lightning's logger. These two handlers are essentially message buffers. This
|
||||
function is called in nemo.utils.__init__.py. These handlers are used in add_filehandlers_to_pl_logger to flush
|
||||
buffered messages to files.
|
||||
"""
|
||||
if not HANDLERS:
|
||||
HANDLERS["memory_err"] = MemoryHandler(-1)
|
||||
HANDLERS["memory_err"].addFilter(lambda record: record.levelno > _logging.INFO)
|
||||
HANDLERS["memory_all"] = MemoryHandler(-1)
|
||||
pl._logger.addHandler(HANDLERS["memory_err"])
|
||||
pl._logger.addHandler(HANDLERS["memory_all"])
|
||||
|
||||
|
||||
def add_filehandlers_to_pl_logger(all_log_file, err_log_file):
|
||||
"""
|
||||
Adds two filehandlers to pytorch_lightning's logger. Called in nemo.utils.exp_manager(). The first filehandler
|
||||
logs all messages to all_log_file while the second filehandler logs all WARNING and higher messages to err_log_file.
|
||||
If "memory_err" and "memory_all" exist in HANDLERS, then those buffers are flushed to err_log_file and all_log_file
|
||||
respectively, and then closed.
|
||||
"""
|
||||
HANDLERS["file"] = _logging.FileHandler(all_log_file)
|
||||
pl._logger.addHandler(HANDLERS["file"])
|
||||
HANDLERS["file_err"] = _logging.FileHandler(err_log_file)
|
||||
HANDLERS["file_err"].addFilter(lambda record: record.levelno > _logging.INFO)
|
||||
pl._logger.addHandler(HANDLERS["file_err"])
|
||||
|
||||
if HANDLERS.get("memory_all", None):
|
||||
HANDLERS["memory_all"].setTarget(HANDLERS["file"])
|
||||
HANDLERS["memory_all"].close()
|
||||
del HANDLERS["memory_all"]
|
||||
if HANDLERS.get("memory_err", None):
|
||||
HANDLERS["memory_err"].setTarget(HANDLERS["file_err"])
|
||||
HANDLERS["memory_err"].close()
|
||||
del HANDLERS["memory_err"]
|
||||
@@ -0,0 +1,39 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import threading
|
||||
|
||||
|
||||
class Singleton(type):
|
||||
""" Implementation of a generic, tread-safe singleton meta-class.
|
||||
Can be used as meta-class, i.e. will create
|
||||
"""
|
||||
|
||||
# List of instances - one per class.
|
||||
__instances = {}
|
||||
# Lock used for accessing the instance.
|
||||
__lock = threading.Lock()
|
||||
|
||||
def __call__(cls, *args, **kwargs):
|
||||
""" Returns singleton instance. A thread safe implementation. """
|
||||
if cls not in cls.__instances:
|
||||
# Enter critical section.
|
||||
with cls.__lock:
|
||||
# Check once again.
|
||||
if cls not in cls.__instances:
|
||||
# Create a new object instance - one per class.
|
||||
cls.__instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
|
||||
# Return the instance.
|
||||
return cls.__instances[cls]
|
||||
@@ -0,0 +1,528 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import copy
|
||||
import os
|
||||
from dataclasses import dataclass, is_dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import pytorch_lightning as pl
|
||||
import wrapt
|
||||
from omegaconf import DictConfig, ListConfig, OmegaConf
|
||||
from omegaconf import errors as omegaconf_errors
|
||||
from packaging import version
|
||||
|
||||
from nemo.utils import logging
|
||||
|
||||
_VAL_TEST_FASTPATH_KEY = 'ds_item'
|
||||
|
||||
|
||||
class ArtifactPathType(Enum):
|
||||
"""
|
||||
ArtifactPathType refers to the type of the path that the artifact is located at.
|
||||
|
||||
LOCAL_PATH: A user local filepath that exists on the file system.
|
||||
TAR_PATH: A (generally flattened) filepath that exists inside of an archive (that may have its own full path).
|
||||
"""
|
||||
|
||||
LOCAL_PATH = 0
|
||||
TAR_PATH = 1
|
||||
|
||||
|
||||
@dataclass(init=False)
|
||||
class ArtifactItem:
|
||||
path: str
|
||||
path_type: ArtifactPathType
|
||||
hashed_path: Optional[str] = None
|
||||
|
||||
|
||||
def resolve_dataset_name_from_cfg(cfg: DictConfig) -> str:
|
||||
"""
|
||||
Parses items of the provided sub-config to find the first potential key that
|
||||
resolves to an existing file or directory.
|
||||
|
||||
# Fast-path Resolution
|
||||
In order to handle cases where we need to resolve items that are not paths, a fastpath
|
||||
key can be provided as defined in the global `_VAL_TEST_FASTPATH_KEY`.
|
||||
|
||||
This key can be used in two ways :
|
||||
|
||||
## _VAL_TEST_FASTPATH_KEY points to another key in the config
|
||||
|
||||
If this _VAL_TEST_FASTPATH_KEY points to another key in this config itself,
|
||||
then we assume we want to loop through the values of that key.
|
||||
|
||||
This allows for any key in the config to become a fastpath key.
|
||||
|
||||
Example:
|
||||
validation_ds:
|
||||
splits: "val"
|
||||
...
|
||||
<_VAL_TEST_FASTPATH_KEY>: "splits" <-- this points to the key name "splits"
|
||||
|
||||
Then we can write the following when overriding in hydra:
|
||||
```python
|
||||
python train_file.py ... \
|
||||
model.validation_ds.splits=[val1, val2, dev1, dev2] ...
|
||||
```
|
||||
|
||||
## _VAL_TEST_FASTPATH_KEY itself acts as the resolved key
|
||||
|
||||
If this _VAL_TEST_FASTPATH_KEY does not point to another key in the config, then
|
||||
it is assumed that the items of this key itself are used for resolution.
|
||||
|
||||
Example:
|
||||
validation_ds:
|
||||
...
|
||||
<_VAL_TEST_FASTPATH_KEY>: "val" <-- this points to the key name "splits"
|
||||
|
||||
Then we can write the following when overriding in hydra:
|
||||
```python
|
||||
python train_file.py ... \
|
||||
model.validation_ds.<_VAL_TEST_FASTPATH_KEY>=[val1, val2, dev1, dev2] ...
|
||||
```
|
||||
|
||||
# IMPORTANT NOTE:
|
||||
It <can> potentially mismatch if there exist more than 2 valid paths, and the
|
||||
first path does *not* resolve the the path of the data file (but does resolve to
|
||||
some other valid path).
|
||||
|
||||
To avoid this side-effect, place the data path as the first item on the config file.
|
||||
|
||||
Args:
|
||||
cfg: DictConfig (Sub-config) that should be parsed.
|
||||
|
||||
Returns:
|
||||
A str representing the `key` of the config which hosts the filepath(s),
|
||||
or None in case path could not be resolved.
|
||||
"""
|
||||
if _VAL_TEST_FASTPATH_KEY in cfg and cfg[_VAL_TEST_FASTPATH_KEY] is not None:
|
||||
fastpath_key = cfg[_VAL_TEST_FASTPATH_KEY]
|
||||
|
||||
if isinstance(fastpath_key, str) and fastpath_key in cfg:
|
||||
return cfg[fastpath_key]
|
||||
else:
|
||||
return _VAL_TEST_FASTPATH_KEY
|
||||
|
||||
for key, value in cfg.items():
|
||||
if type(value) in [list, tuple, ListConfig]:
|
||||
# Count the number of valid paths in the list
|
||||
values_are_paths = 0
|
||||
for val_i in value:
|
||||
val_i = str(val_i)
|
||||
|
||||
if os.path.exists(val_i) or os.path.isdir(val_i):
|
||||
values_are_paths += 1
|
||||
else:
|
||||
# reset counter and break inner loop
|
||||
break
|
||||
|
||||
if values_are_paths == len(value):
|
||||
return key
|
||||
|
||||
else:
|
||||
if os.path.exists(str(value)) or os.path.isdir(str(value)):
|
||||
return key
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def parse_dataset_as_name(name: str) -> str:
|
||||
"""
|
||||
Constructs a valid prefix-name from a provided file path.
|
||||
|
||||
Args:
|
||||
name: str path to some valid data/manifest file or a python object that
|
||||
will be used as a name for the data loader (via str() cast).
|
||||
|
||||
Returns:
|
||||
str prefix used to identify uniquely this data/manifest file.
|
||||
"""
|
||||
if os.path.exists(str(name)) or os.path.isdir(str(name)):
|
||||
name = Path(name).stem
|
||||
else:
|
||||
name = str(name)
|
||||
|
||||
# cleanup name
|
||||
name = name.replace('-', '_')
|
||||
|
||||
if 'manifest' in name:
|
||||
name = name.replace('manifest', '')
|
||||
|
||||
if 'dataset' in name:
|
||||
name = name.replace('dataset', '')
|
||||
|
||||
if '_' != name[-1]:
|
||||
name = name + '_'
|
||||
|
||||
return name
|
||||
|
||||
|
||||
def unique_names_check(name_list: Optional[List[str]]):
|
||||
"""
|
||||
Performs a uniqueness check on the name list resolved, so that it can warn users
|
||||
about non-unique keys.
|
||||
|
||||
Args:
|
||||
name_list: List of strings resolved for data loaders.
|
||||
"""
|
||||
if name_list is None:
|
||||
return
|
||||
|
||||
# Name uniqueness checks
|
||||
names = set()
|
||||
for name in name_list:
|
||||
if name in names:
|
||||
logging.warning(
|
||||
"Name resolution has found more than one data loader having the same name !\n"
|
||||
"In such cases, logs will nor be properly generated. "
|
||||
"Please rename the item to have unique names.\n"
|
||||
f"Resolved name : {name}"
|
||||
)
|
||||
else:
|
||||
names.add(name) # we need just hash key check, value is just a placeholder
|
||||
|
||||
|
||||
def resolve_validation_dataloaders(model: 'ModelPT'):
|
||||
"""
|
||||
Helper method that operates on the ModelPT class to automatically support
|
||||
multiple dataloaders for the validation set.
|
||||
|
||||
It does so by first resolving the path to one/more data files via `resolve_dataset_name_from_cfg()`.
|
||||
If this resolution fails, it assumes the data loader is prepared to manually support / not support
|
||||
multiple data loaders and simply calls the appropriate setup method.
|
||||
|
||||
If resolution succeeds:
|
||||
Checks if provided path is to a single file or a list of files.
|
||||
If a single file is provided, simply tags that file as such and loads it via the setup method.
|
||||
If multiple files are provided:
|
||||
Inject a new manifest path at index "i" into the resolved key.
|
||||
Calls the appropriate setup method to set the data loader.
|
||||
Collects the initialized data loader in a list and preserves it.
|
||||
Once all data loaders are processed, assigns the list of loaded loaders to the ModelPT.
|
||||
Finally assigns a list of unique names resolved from the file paths to the ModelPT.
|
||||
|
||||
Args:
|
||||
model: ModelPT subclass, which requires >=1 Validation Dataloaders to be setup.
|
||||
"""
|
||||
cfg = copy.deepcopy(model._cfg)
|
||||
dataloaders = []
|
||||
|
||||
# process val_loss_idx
|
||||
if 'val_dl_idx' in cfg.validation_ds:
|
||||
cfg = OmegaConf.to_container(cfg)
|
||||
val_dl_idx = cfg['validation_ds'].pop('val_dl_idx')
|
||||
cfg = OmegaConf.create(cfg)
|
||||
else:
|
||||
val_dl_idx = 0
|
||||
|
||||
# Set val_loss_idx
|
||||
model._val_dl_idx = val_dl_idx
|
||||
|
||||
ds_key = resolve_dataset_name_from_cfg(cfg.validation_ds)
|
||||
|
||||
if ds_key is None:
|
||||
logging.debug(
|
||||
"Could not resolve file path from provided config - {}. "
|
||||
"Disabling support for multi-dataloaders.".format(cfg.validation_ds)
|
||||
)
|
||||
|
||||
model.setup_validation_data(cfg.validation_ds)
|
||||
return
|
||||
|
||||
ds_values = cfg.validation_ds[ds_key]
|
||||
|
||||
if isinstance(ds_values, (list, tuple, ListConfig)):
|
||||
|
||||
for ds_value in ds_values:
|
||||
cfg.validation_ds[ds_key] = ds_value
|
||||
model.setup_validation_data(cfg.validation_ds)
|
||||
dataloaders.append(model._validation_dl)
|
||||
|
||||
model._validation_dl = dataloaders
|
||||
model._validation_names = [parse_dataset_as_name(ds) for ds in ds_values]
|
||||
|
||||
unique_names_check(name_list=model._validation_names)
|
||||
return
|
||||
|
||||
else:
|
||||
model.setup_validation_data(cfg.validation_ds)
|
||||
model._validation_names = [parse_dataset_as_name(ds_values)]
|
||||
|
||||
unique_names_check(name_list=model._validation_names)
|
||||
|
||||
|
||||
def resolve_test_dataloaders(model: 'ModelPT'):
|
||||
"""
|
||||
Helper method that operates on the ModelPT class to automatically support
|
||||
multiple dataloaders for the test set.
|
||||
|
||||
It does so by first resolving the path to one/more data files via `resolve_dataset_name_from_cfg()`.
|
||||
If this resolution fails, it assumes the data loader is prepared to manually support / not support
|
||||
multiple data loaders and simply calls the appropriate setup method.
|
||||
|
||||
If resolution succeeds:
|
||||
Checks if provided path is to a single file or a list of files.
|
||||
If a single file is provided, simply tags that file as such and loads it via the setup method.
|
||||
If multiple files are provided:
|
||||
Inject a new manifest path at index "i" into the resolved key.
|
||||
Calls the appropriate setup method to set the data loader.
|
||||
Collects the initialized data loader in a list and preserves it.
|
||||
Once all data loaders are processed, assigns the list of loaded loaders to the ModelPT.
|
||||
Finally assigns a list of unique names resolved from the file paths to the ModelPT.
|
||||
|
||||
Args:
|
||||
model: ModelPT subclass, which requires >=1 Test Dataloaders to be setup.
|
||||
"""
|
||||
cfg = copy.deepcopy(model._cfg)
|
||||
dataloaders = []
|
||||
|
||||
# process test_loss_idx
|
||||
if 'test_dl_idx' in cfg.test_ds:
|
||||
cfg = OmegaConf.to_container(cfg)
|
||||
test_dl_idx = cfg['test_ds'].pop('test_dl_idx')
|
||||
cfg = OmegaConf.create(cfg)
|
||||
else:
|
||||
test_dl_idx = 0
|
||||
|
||||
# Set val_loss_idx
|
||||
model._test_dl_idx = test_dl_idx
|
||||
|
||||
ds_key = resolve_dataset_name_from_cfg(cfg.test_ds)
|
||||
|
||||
if ds_key is None:
|
||||
logging.debug(
|
||||
"Could not resolve file path from provided config - {}. "
|
||||
"Disabling support for multi-dataloaders.".format(cfg.test_ds)
|
||||
)
|
||||
|
||||
model.setup_test_data(cfg.test_ds)
|
||||
return
|
||||
|
||||
ds_values = cfg.test_ds[ds_key]
|
||||
|
||||
if isinstance(ds_values, (list, tuple, ListConfig)):
|
||||
|
||||
for ds_value in ds_values:
|
||||
cfg.test_ds[ds_key] = ds_value
|
||||
model.setup_test_data(cfg.test_ds)
|
||||
dataloaders.append(model._test_dl)
|
||||
|
||||
model._test_dl = dataloaders
|
||||
model._test_names = [parse_dataset_as_name(ds) for ds in ds_values]
|
||||
|
||||
unique_names_check(name_list=model._test_names)
|
||||
return
|
||||
|
||||
else:
|
||||
model.setup_test_data(cfg.test_ds)
|
||||
model._test_names = [parse_dataset_as_name(ds_values)]
|
||||
|
||||
unique_names_check(name_list=model._test_names)
|
||||
|
||||
|
||||
@wrapt.decorator
|
||||
def wrap_training_step(wrapped, instance: pl.LightningModule, args, kwargs):
|
||||
output_dict = wrapped(*args, **kwargs)
|
||||
|
||||
if isinstance(output_dict, dict) and output_dict is not None and 'log' in output_dict:
|
||||
log_dict = output_dict.pop('log')
|
||||
instance.log_dict(log_dict, on_step=True)
|
||||
|
||||
return output_dict
|
||||
|
||||
|
||||
def convert_model_config_to_dict_config(cfg: Union[DictConfig, 'NemoConfig']) -> DictConfig:
|
||||
"""
|
||||
Converts its input into a standard DictConfig.
|
||||
Possible input values are:
|
||||
- DictConfig
|
||||
- A dataclass which is a subclass of NemoConfig
|
||||
|
||||
Args:
|
||||
cfg: A dict-like object.
|
||||
|
||||
Returns:
|
||||
The equivalent DictConfig
|
||||
"""
|
||||
if not isinstance(cfg, (OmegaConf, DictConfig)) and is_dataclass(cfg):
|
||||
cfg = OmegaConf.structured(cfg)
|
||||
|
||||
if not isinstance(cfg, DictConfig):
|
||||
raise ValueError(f"cfg constructor argument must be of type DictConfig/dict but got {type(cfg)} instead.")
|
||||
|
||||
config = OmegaConf.to_container(cfg, resolve=True)
|
||||
config = OmegaConf.create(config)
|
||||
return config
|
||||
|
||||
|
||||
def _convert_config(cfg: OmegaConf):
|
||||
""" Recursive function convertint the configuration from old hydra format to the new one. """
|
||||
|
||||
# Get rid of cls -> _target_.
|
||||
if 'cls' in cfg and '_target_' not in cfg:
|
||||
cfg._target_ = cfg.pop('cls')
|
||||
|
||||
# Get rid of params.
|
||||
if 'params' in cfg:
|
||||
params = cfg.pop('params')
|
||||
for param_key, param_val in params.items():
|
||||
cfg[param_key] = param_val
|
||||
|
||||
# Recursion.
|
||||
try:
|
||||
for _, sub_cfg in cfg.items():
|
||||
if isinstance(sub_cfg, DictConfig):
|
||||
_convert_config(sub_cfg)
|
||||
except omegaconf_errors.OmegaConfBaseException as e:
|
||||
logging.warning(f"Skipped conversion for config/subconfig:\n{cfg}\n Reason: {e}.")
|
||||
|
||||
|
||||
def maybe_update_config_version(cfg: DictConfig):
|
||||
"""
|
||||
Recursively convert Hydra 0.x configs to Hydra 1.x configs.
|
||||
|
||||
Changes include:
|
||||
- `cls` -> `_target_`.
|
||||
- `params` -> drop params and shift all arguments to parent.
|
||||
- `target` -> `_target_` cannot be performed due to ModelPT injecting `target` inside class.
|
||||
|
||||
Args:
|
||||
cfg: Any Hydra compatible DictConfig
|
||||
|
||||
Returns:
|
||||
An updated DictConfig that conforms to Hydra 1.x format.
|
||||
"""
|
||||
if cfg is not None and not isinstance(cfg, DictConfig):
|
||||
try:
|
||||
temp_cfg = OmegaConf.create(cfg)
|
||||
cfg = temp_cfg
|
||||
except omegaconf_errors.OmegaConfBaseException:
|
||||
# Cannot be cast to DictConfig, skip updating.
|
||||
return cfg
|
||||
|
||||
# Make a copy of model config.
|
||||
cfg = copy.deepcopy(cfg)
|
||||
OmegaConf.set_struct(cfg, False)
|
||||
|
||||
# Convert config.
|
||||
_convert_config(cfg)
|
||||
|
||||
# Update model config.
|
||||
OmegaConf.set_struct(cfg, True)
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
def import_class_by_path(path: str):
|
||||
"""
|
||||
Recursive import of class by path string.
|
||||
"""
|
||||
paths = path.split('.')
|
||||
path = ".".join(paths[:-1])
|
||||
class_name = paths[-1]
|
||||
mod = __import__(path, fromlist=[class_name])
|
||||
mod = getattr(mod, class_name)
|
||||
return mod
|
||||
|
||||
|
||||
def resolve_subclass_pretrained_model_info(base_class) -> List['PretrainedModelInfo']:
|
||||
"""
|
||||
Recursively traverses the inheritance graph of subclasses to extract all pretrained model info.
|
||||
First constructs a set of unique pretrained model info by performing DFS over the inheritance graph.
|
||||
All model info belonging to the same class is added together.
|
||||
|
||||
Args:
|
||||
base_class: The root class, whose subclass graph will be traversed.
|
||||
|
||||
Returns:
|
||||
A list of unique pretrained model infos belonging to all of the inherited subclasses of
|
||||
this baseclass.
|
||||
"""
|
||||
list_of_models = set()
|
||||
|
||||
def recursive_subclass_walk(cls):
|
||||
for subclass in cls.__subclasses__():
|
||||
# step into its immediate subclass
|
||||
recursive_subclass_walk(subclass)
|
||||
|
||||
subclass_models = subclass.list_available_models()
|
||||
|
||||
if subclass_models is not None and len(subclass_models) > 0:
|
||||
# Inject subclass info into pretrained model info
|
||||
# if not already overriden by subclass
|
||||
for model_info in subclass_models:
|
||||
# If subclass manually injects class_, dont override.
|
||||
if model_info.class_ is None:
|
||||
model_info.class_ = subclass
|
||||
|
||||
for model_info in subclass_models:
|
||||
list_of_models.add(model_info)
|
||||
|
||||
recursive_subclass_walk(base_class)
|
||||
|
||||
list_of_models = list(sorted(list_of_models))
|
||||
return list_of_models
|
||||
|
||||
|
||||
def check_lib_version(lib_name: str, checked_version: str, operator) -> (Optional[bool], str):
|
||||
"""
|
||||
Checks if a library is installed, and if it is, checks the operator(lib.__version__, checked_version) as a result.
|
||||
This bool result along with a string analysis of result is returned.
|
||||
|
||||
If the library is not installed at all, then returns None instead, along with a string explaining
|
||||
that the library is not installed
|
||||
|
||||
Args:
|
||||
lib_name: lower case str name of the library that must be imported.
|
||||
checked_version: semver string that is compared against lib.__version__.
|
||||
operator: binary callable function func(a, b) -> bool; that compares lib.__version__ against version in
|
||||
some manner. Must return a boolean.
|
||||
|
||||
Returns:
|
||||
A tuple of results:
|
||||
- Bool or None. Bool if the library could be imported, and the result of
|
||||
operator(lib.__version__, checked_version) or False if __version__ is not implemented in lib.
|
||||
None is passed if the library is not installed at all.
|
||||
- A string analysis of the check.
|
||||
"""
|
||||
try:
|
||||
mod = __import__(lib_name)
|
||||
|
||||
if hasattr(mod, '__version__'):
|
||||
lib_ver = version.Version(mod.__version__)
|
||||
match_ver = version.Version(checked_version)
|
||||
|
||||
if operator(lib_ver, match_ver):
|
||||
msg = f"Lib {lib_name} version is satisfied !"
|
||||
return True, msg
|
||||
else:
|
||||
msg = (
|
||||
f"Lib {lib_name} version ({lib_ver}) is not {operator.__name__} than required version {checked_version}.\n"
|
||||
f"Please upgrade the lib using either pip or conda to the latest version."
|
||||
)
|
||||
return False, msg
|
||||
else:
|
||||
msg = (
|
||||
f"Lib {lib_name} does not implement __version__ in its init file. "
|
||||
f"Could not check version compatibility."
|
||||
)
|
||||
return False, msg
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
pass
|
||||
|
||||
msg = f"Lib {lib_name} has not been installed. Please use pip or conda to install this package."
|
||||
return None, msg
|
||||
@@ -0,0 +1,421 @@
|
||||
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import enum
|
||||
import logging as _logging
|
||||
import sys
|
||||
import threading
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
from logging.handlers import MemoryHandler
|
||||
|
||||
from nemo.constants import NEMO_ENV_VARNAME_REDIRECT_LOGS_TO_STDERR, NEMO_ENV_VARNAME_TESTING
|
||||
from nemo.utils.env_var_parsing import get_envbool
|
||||
from nemo.utils.formatters.base import BaseNeMoFormatter, DebugNeMoFormatter
|
||||
from nemo.utils.get_rank import is_global_rank_zero
|
||||
from nemo.utils.metaclasses import Singleton
|
||||
|
||||
__all__ = ["Logger", "LogMode"]
|
||||
|
||||
|
||||
class LogMode(enum.IntEnum):
|
||||
EACH = 0 # Log the message each time
|
||||
ONCE = 1 # Log the message only once. The same message will not be logged again.
|
||||
|
||||
|
||||
class Logger(metaclass=Singleton):
|
||||
|
||||
# Level 0
|
||||
NOTSET = _logging.NOTSET
|
||||
|
||||
# Level 10
|
||||
DEBUG = _logging.DEBUG
|
||||
|
||||
# Level 20
|
||||
INFO = _logging.INFO
|
||||
|
||||
# Level 30
|
||||
WARNING = _logging.WARNING
|
||||
|
||||
# Level 40
|
||||
ERROR = _logging.ERROR
|
||||
|
||||
# Level 50
|
||||
CRITICAL = _logging.CRITICAL
|
||||
|
||||
_level_names = {
|
||||
0: "NOTSET",
|
||||
10: "DEBUG",
|
||||
20: "INFO",
|
||||
30: "WARNING",
|
||||
40: "ERROR",
|
||||
50: "CRITICAL",
|
||||
}
|
||||
|
||||
def __init__(self, capture_warnings=True):
|
||||
|
||||
self._logger = None
|
||||
# Multi-GPU runs run in separate processes, thread locks shouldn't be needed
|
||||
self._logger_lock = threading.Lock()
|
||||
self._handlers = dict()
|
||||
self.old_warnings_showwarning = None
|
||||
self._define_logger(capture_warnings)
|
||||
self.once_logged = set()
|
||||
self.rank = 0 if is_global_rank_zero() else "UNK"
|
||||
|
||||
def _define_logger(self, capture_warnings=True):
|
||||
""" Creates the logger if not already created. Called in init"""
|
||||
|
||||
# Use double-checked locking to avoid taking lock unnecessarily.
|
||||
if self._logger is not None:
|
||||
return self._logger
|
||||
|
||||
with self._logger_lock:
|
||||
try:
|
||||
self._logger = _logging.getLogger("nemo_logger")
|
||||
# By default, silence all loggers except the logger for rank 0
|
||||
self.remove_stream_handlers()
|
||||
# If NEMO_TESTING is set, add a streamhandler to all ranks
|
||||
if get_envbool(NEMO_ENV_VARNAME_TESTING, False):
|
||||
old_factory = _logging.getLogRecordFactory()
|
||||
|
||||
def record_factory(*args, **kwargs):
|
||||
record = old_factory(*args, **kwargs)
|
||||
record.rank = self.rank
|
||||
return record
|
||||
|
||||
_logging.setLogRecordFactory(record_factory)
|
||||
self.add_stream_handlers(formatter=DebugNeMoFormatter)
|
||||
elif is_global_rank_zero():
|
||||
self.add_stream_handlers()
|
||||
|
||||
# Add memoryhandlers, essentially buffers. They are used to save messages that we will flush to file
|
||||
# once the appropriate file handlers are added.
|
||||
if is_global_rank_zero():
|
||||
# Add a memoryhandler for error messages. Only logged on rank 0
|
||||
self._handlers["memory_err"] = MemoryHandler(-1)
|
||||
self._handlers["memory_err"].addFilter(lambda record: record.levelno > _logging.INFO)
|
||||
formatter = BaseNeMoFormatter
|
||||
self._handlers["memory_err"].setFormatter(formatter())
|
||||
self._logger.addHandler(self._handlers["memory_err"])
|
||||
# Add a memoryhandler for all messages on all ranks
|
||||
self._handlers["memory_all"] = MemoryHandler(-1)
|
||||
formatter = BaseNeMoFormatter
|
||||
self._handlers["memory_all"].setFormatter(formatter())
|
||||
self._logger.addHandler(self._handlers["memory_all"])
|
||||
|
||||
finally:
|
||||
level = Logger.INFO
|
||||
if get_envbool(NEMO_ENV_VARNAME_TESTING, False):
|
||||
level = Logger.DEBUG
|
||||
self.set_verbosity(verbosity_level=level)
|
||||
self.captureWarnings(capture_warnings)
|
||||
|
||||
self._logger.propagate = False
|
||||
|
||||
def remove_stream_handlers(self):
|
||||
""" Removes StreamHandler that log to stdout and stderr from the logger."""
|
||||
if self._logger is None:
|
||||
raise RuntimeError("Impossible to set handlers if the Logger is not predefined")
|
||||
|
||||
# ======== Remove Handler if already existing ========
|
||||
|
||||
try:
|
||||
self._logger.removeHandler(self._handlers["stream_stdout"])
|
||||
del self._handlers["stream_stdout"]
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
try:
|
||||
self._logger.removeHandler(self._handlers["stream_stderr"])
|
||||
del self._handlers["stream_stderr"]
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
def add_stream_handlers(self, formatter=BaseNeMoFormatter):
|
||||
"""Add StreamHandler that log to stdout and stderr to the logger. INFO and lower logs are streamed to stdout
|
||||
while WARNING and higher are streamed to stderr. If the NEMO_ENV_VARNAME_REDIRECT_LOGS_TO_STDERR environment
|
||||
variable is set, all logs are sent to stderr instead.
|
||||
"""
|
||||
if self._logger is None:
|
||||
raise RuntimeError("Impossible to set handlers if the Logger is not predefined")
|
||||
|
||||
# Add the output handler.
|
||||
if get_envbool(NEMO_ENV_VARNAME_REDIRECT_LOGS_TO_STDERR, False):
|
||||
self._handlers["stream_stdout"] = _logging.StreamHandler(sys.stderr)
|
||||
|
||||
else:
|
||||
self._handlers["stream_stdout"] = _logging.StreamHandler(sys.stdout)
|
||||
self._handlers["stream_stdout"].addFilter(lambda record: record.levelno <= _logging.INFO)
|
||||
|
||||
self._handlers["stream_stderr"] = _logging.StreamHandler(sys.stderr)
|
||||
self._handlers["stream_stderr"].addFilter(lambda record: record.levelno > _logging.INFO)
|
||||
|
||||
self._handlers["stream_stdout"].setFormatter(formatter())
|
||||
self._logger.addHandler(self._handlers["stream_stdout"])
|
||||
|
||||
try:
|
||||
self._handlers["stream_stderr"].setFormatter(formatter())
|
||||
self._logger.addHandler(self._handlers["stream_stderr"])
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
def reset_stream_handler(self, formatter=BaseNeMoFormatter):
|
||||
"""Removes then adds stream handlers."""
|
||||
self.remove_stream_handlers()
|
||||
self.add_stream_handlers(formatter=formatter)
|
||||
|
||||
def add_file_handler(self, log_file):
|
||||
"""Add a FileHandler to logger that logs all messages to a file. If the logger had a MemoryHandler at
|
||||
self._handlers["memory_all"], those buffered messages are flushed to the new file, and the MemoryHandler is
|
||||
closed."""
|
||||
if self._logger is None:
|
||||
raise RuntimeError("Impossible to set handlers if the Logger is not predefined")
|
||||
|
||||
self._handlers["file"] = _logging.FileHandler(log_file)
|
||||
formatter = BaseNeMoFormatter
|
||||
self._handlers["file"].setFormatter(formatter())
|
||||
self._logger.addHandler(self._handlers["file"])
|
||||
|
||||
if self._handlers.get("memory_all", None):
|
||||
self._handlers["memory_all"].setTarget(self._handlers["file"])
|
||||
self._handlers["memory_all"].close() # flush and remove
|
||||
del self._handlers["memory_all"]
|
||||
|
||||
def add_err_file_handler(self, log_file):
|
||||
"""Add a FileHandler to logger that logs all WARNING and higher messages to a file. If the logger had a
|
||||
MemoryHandler at self._handlers["memory_err"], those buffered messages are flushed to the new file, and the
|
||||
MemoryHandler is closed."""
|
||||
if self._logger is None:
|
||||
raise RuntimeError("Impossible to set handlers if the Logger is not predefined")
|
||||
|
||||
self._handlers["file_err"] = _logging.FileHandler(log_file)
|
||||
self._handlers["file_err"].addFilter(lambda record: record.levelno > _logging.INFO)
|
||||
|
||||
formatter = BaseNeMoFormatter
|
||||
self._handlers["file_err"].setFormatter(formatter())
|
||||
self._logger.addHandler(self._handlers["file_err"])
|
||||
|
||||
if self._handlers.get("memory_err", None):
|
||||
self._handlers["memory_err"].setTarget(self._handlers["file_err"])
|
||||
self._handlers["memory_err"].close() # flush and remove
|
||||
del self._handlers["memory_err"]
|
||||
|
||||
def getEffectiveLevel(self):
|
||||
"""Return how much logging output will be produced."""
|
||||
if self._logger is not None:
|
||||
return self._logger.getEffectiveLevel()
|
||||
|
||||
def get_verbosity(self):
|
||||
"""See getEffectiveLevel"""
|
||||
return self.getEffectiveLevel()
|
||||
|
||||
def setLevel(self, verbosity_level):
|
||||
"""Sets the threshold for what messages will be logged."""
|
||||
if self._logger is not None:
|
||||
self._logger.setLevel(verbosity_level)
|
||||
|
||||
for handler in self._logger.handlers:
|
||||
handler.setLevel(verbosity_level)
|
||||
|
||||
def set_verbosity(self, verbosity_level):
|
||||
"""See setLevel"""
|
||||
self.setLevel(verbosity_level)
|
||||
|
||||
@contextmanager
|
||||
def patch_stderr_handler(self, stream):
|
||||
""" Sends messages that should log to stderr to stream instead. Useful for unittests """
|
||||
if self._logger is not None:
|
||||
try:
|
||||
old_stream = self._handlers["stream_stderr"].stream
|
||||
if old_stream is None:
|
||||
raise ValueError
|
||||
|
||||
# Port backwards set_stream() from python 3.7
|
||||
self._handlers["stream_stderr"].acquire()
|
||||
try:
|
||||
self._handlers["stream_stderr"].flush()
|
||||
self._handlers["stream_stderr"].stream = stream
|
||||
finally:
|
||||
self._handlers["stream_stderr"].release()
|
||||
|
||||
yield stream
|
||||
except (KeyError, ValueError):
|
||||
raise RuntimeError("Impossible to patch logging handlers if handler does not exist")
|
||||
finally:
|
||||
# Port backwards set_stream() from python 3.7
|
||||
self._handlers["stream_stderr"].acquire()
|
||||
try:
|
||||
self._handlers["stream_stderr"].flush()
|
||||
self._handlers["stream_stderr"].stream = old_stream
|
||||
finally:
|
||||
self._handlers["stream_stderr"].release()
|
||||
|
||||
else:
|
||||
raise RuntimeError("Impossible to patch logging handlers if handler does not exist")
|
||||
|
||||
@contextmanager
|
||||
def patch_stdout_handler(self, stream):
|
||||
""" Sends messages that should log to stdout to stream instead. Useful for unittests """
|
||||
if self._logger is not None:
|
||||
try:
|
||||
old_stream = self._handlers["stream_stdout"].stream
|
||||
if old_stream is None:
|
||||
raise ValueError
|
||||
|
||||
# Port backwards set_stream() from python 3.7
|
||||
self._handlers["stream_stdout"].acquire()
|
||||
try:
|
||||
self._handlers["stream_stdout"].flush()
|
||||
self._handlers["stream_stdout"].stream = stream
|
||||
finally:
|
||||
self._handlers["stream_stdout"].release()
|
||||
|
||||
yield stream
|
||||
except (KeyError, ValueError):
|
||||
raise RuntimeError("Impossible to patch logging handlers if handler does not exist")
|
||||
finally:
|
||||
# Port backwards set_stream() from python 3.7
|
||||
self._handlers["stream_stdout"].acquire()
|
||||
try:
|
||||
self._handlers["stream_stdout"].flush()
|
||||
self._handlers["stream_stdout"].stream = old_stream
|
||||
finally:
|
||||
self._handlers["stream_stdout"].release()
|
||||
|
||||
else:
|
||||
raise RuntimeError("Impossible to patch logging handlers if handler does not exist")
|
||||
|
||||
@contextmanager
|
||||
def temp_verbosity(self, verbosity_level):
|
||||
"""Sets the a temporary threshold for what messages will be logged."""
|
||||
|
||||
if self._logger is not None:
|
||||
|
||||
old_verbosity = self.get_verbosity()
|
||||
|
||||
try:
|
||||
self.set_verbosity(verbosity_level)
|
||||
yield
|
||||
|
||||
finally:
|
||||
self.set_verbosity(old_verbosity)
|
||||
|
||||
else:
|
||||
try:
|
||||
yield
|
||||
|
||||
finally:
|
||||
pass
|
||||
|
||||
def captureWarnings(self, capture):
|
||||
"""
|
||||
If capture is true, redirect all warnings to the logging package.
|
||||
If capture is False, ensure that warnings are not redirected to logging
|
||||
but to their original destinations.
|
||||
"""
|
||||
|
||||
if self._logger is not None:
|
||||
|
||||
if capture and self.old_warnings_showwarning is None:
|
||||
# Backup Method
|
||||
self.old_warnings_showwarning = warnings.showwarning
|
||||
warnings.showwarning = self._showwarning
|
||||
|
||||
elif not capture and self.old_warnings_showwarning is not None:
|
||||
# Restore Method
|
||||
warnings.showwarning = self.old_warnings_showwarning
|
||||
self.old_warnings_showwarning = None
|
||||
|
||||
def _showwarning(self, message, category, filename, lineno, file=None, line=None):
|
||||
"""
|
||||
Implementation of showwarnings which redirects to logging.
|
||||
It will call warnings.formatwarning and will log the resulting string
|
||||
with level logging.WARNING.
|
||||
"""
|
||||
s = warnings.formatwarning(message, category, filename, lineno, line)
|
||||
self.warning("%s", s)
|
||||
|
||||
def _logged_once(self, msg, mode):
|
||||
PREFIX_LEN = 12
|
||||
if mode == LogMode.ONCE:
|
||||
if msg[PREFIX_LEN:] in self.once_logged:
|
||||
return True
|
||||
self.once_logged.add(msg[PREFIX_LEN:])
|
||||
return False
|
||||
|
||||
def debug(self, msg, *args, mode=LogMode.EACH, **kwargs):
|
||||
"""
|
||||
Log 'msg % args' with severity 'DEBUG'.
|
||||
|
||||
To pass exception information, use the keyword argument exc_info with
|
||||
a true value, e.g.
|
||||
|
||||
logger.debug("Houston, we have a %s", "thorny problem", exc_info=1)
|
||||
"""
|
||||
if self._logger is not None and self._logger.isEnabledFor(Logger.DEBUG) and not self._logged_once(msg, mode):
|
||||
self._logger._log(Logger.DEBUG, msg, args, **kwargs)
|
||||
|
||||
def info(self, msg, *args, mode=LogMode.EACH, **kwargs):
|
||||
"""
|
||||
Log 'msg % args' with severity 'INFO'.
|
||||
|
||||
To pass exception information, use the keyword argument exc_info with
|
||||
a true value, e.g.
|
||||
|
||||
logger.info("Houston, we have a %s", "interesting problem", exc_info=1)
|
||||
"""
|
||||
if self._logger is not None and self._logger.isEnabledFor(Logger.INFO) and not self._logged_once(msg, mode):
|
||||
self._logger._log(Logger.INFO, msg, args, **kwargs)
|
||||
|
||||
def warning(self, msg, *args, mode=LogMode.EACH, **kwargs):
|
||||
"""
|
||||
Log 'msg % args' with severity 'WARNING'.
|
||||
|
||||
To pass exception information, use the keyword argument exc_info with
|
||||
a true value, e.g.
|
||||
|
||||
logger.warning("Houston, we have a %s", "bit of a problem", exc_info=1)
|
||||
"""
|
||||
if self._logger is not None and self._logger.isEnabledFor(Logger.WARNING) and not self._logged_once(msg, mode):
|
||||
self._logger._log(Logger.WARNING, msg, args, **kwargs)
|
||||
|
||||
def error(self, msg, *args, mode=LogMode.EACH, **kwargs):
|
||||
"""
|
||||
Log 'msg % args' with severity 'ERROR'.
|
||||
|
||||
To pass exception information, use the keyword argument exc_info with
|
||||
a true value, e.g.
|
||||
|
||||
logger.error("Houston, we have a %s", "major problem", exc_info=1)
|
||||
"""
|
||||
if self._logger is not None and self._logger.isEnabledFor(Logger.ERROR) and not self._logged_once(msg, mode):
|
||||
self._logger._log(Logger.ERROR, msg, args, **kwargs)
|
||||
|
||||
def critical(self, msg, *args, mode=LogMode.EACH, **kwargs):
|
||||
"""
|
||||
Log 'msg % args' with severity 'CRITICAL'.
|
||||
|
||||
To pass exception information, use the keyword argument exc_info with
|
||||
a true value, e.g.
|
||||
|
||||
logger.critical("Houston, we have a %s", "major disaster", exc_info=1)
|
||||
"""
|
||||
if (
|
||||
self._logger is not None
|
||||
and self._logger.isEnabledFor(Logger.CRITICAL)
|
||||
and not self._logged_once(msg, mode)
|
||||
):
|
||||
self._logger._log(Logger.CRITICAL, msg, args, **kwargs)
|
||||
Reference in New Issue
Block a user