From 564336967f29c6662b92e0ce91d8b7ded3773833 Mon Sep 17 00:00:00 2001 From: rahul-tuli Date: Tue, 17 Oct 2023 09:19:47 -0400 Subject: [PATCH] Move logger.py to core Add moved logger file to sparseml/pytorch/utils/__init__.py for backwards compatibility --- src/sparseml/core/__init__.py | 1 + src/sparseml/{pytorch/utils => core}/logger.py | 0 .../sparsification/pruning/modifier_pruning_base.py | 2 +- .../sparsification/pruning/modifier_pruning_layer.py | 8 ++++++-- .../sparsification/pruning/modifier_pruning_mfac.py | 3 +-- .../sparsification/pruning/modifier_pruning_obs.py | 3 +-- src/sparseml/pytorch/torchvision/train.py | 12 ++++++------ src/sparseml/pytorch/utils/__init__.py | 2 +- .../pytorch/utils/log_sparsification_info.py | 2 +- src/sparseml/pytorch/utils/module.py | 2 +- src/sparseml/yolov8/trainers.py | 8 ++++++-- .../sparseml/{pytorch/utils => core}/test_logger.py | 2 +- 12 files changed, 26 insertions(+), 19 deletions(-) rename src/sparseml/{pytorch/utils => core}/logger.py (100%) rename tests/sparseml/{pytorch/utils => core}/test_logger.py (98%) diff --git a/src/sparseml/core/__init__.py b/src/sparseml/core/__init__.py index 3a6fc72cb57..ebae013b702 100644 --- a/src/sparseml/core/__init__.py +++ b/src/sparseml/core/__init__.py @@ -20,6 +20,7 @@ from .framework import * from .framework_object import * from .lifecycle import * +from .logger import * from .model import * from .modifier import * from .optimizer import * diff --git a/src/sparseml/pytorch/utils/logger.py b/src/sparseml/core/logger.py similarity index 100% rename from src/sparseml/pytorch/utils/logger.py rename to src/sparseml/core/logger.py diff --git a/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_base.py b/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_base.py index 87a838251c7..72f0f3c78ca 100644 --- a/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_base.py +++ b/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_base.py @@ -36,12 +36,12 @@ from sparseml.pytorch.sparsification.pruning.mask_params import ModuleParamPruningMask from sparseml.pytorch.sparsification.pruning.scorer import PruningParamsScorer from sparseml.pytorch.utils import ( + LoggerManager, NamedLayerParam, get_named_layers_and_params_by_regex, get_prunable_layers, tensor_sparsity, ) -from sparseml.pytorch.utils.logger import LoggerManager from sparseml.sparsification import SparsificationTypes from sparseml.utils import ( ALL_PRUNABLE_TOKEN, diff --git a/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_layer.py b/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_layer.py index 3594c15f73d..ed66dd587c5 100644 --- a/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_layer.py +++ b/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_layer.py @@ -30,8 +30,12 @@ ScheduledModifier, ScheduledUpdateModifier, ) -from sparseml.pytorch.utils import get_layer, get_prunable_layers, replace_layer -from sparseml.pytorch.utils.logger import BaseLogger +from sparseml.pytorch.utils import ( + BaseLogger, + get_layer, + get_prunable_layers, + replace_layer, +) from sparseml.sparsification import SparsificationTypes from sparseml.utils import ALL_PRUNABLE_TOKEN, ALL_TOKEN, validate_str_iterable diff --git a/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_mfac.py b/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_mfac.py index 2150364604f..c6c23c4c255 100644 --- a/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_mfac.py +++ b/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_mfac.py @@ -39,8 +39,7 @@ BaseGradualPruningModifier, ) from sparseml.pytorch.sparsification.pruning.scorer import PruningParamsGradScorer -from sparseml.pytorch.utils import GradSampler -from sparseml.pytorch.utils.logger import BaseLogger +from sparseml.pytorch.utils import BaseLogger, GradSampler __all__ = [ diff --git a/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_obs.py b/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_obs.py index c82cb36a206..fc5d351798f 100644 --- a/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_obs.py +++ b/src/sparseml/pytorch/sparsification/pruning/modifier_pruning_obs.py @@ -37,8 +37,7 @@ BaseGradualPruningModifier, ) from sparseml.pytorch.sparsification.pruning.scorer import PruningParamsGradScorer -from sparseml.pytorch.utils import MEMORY_BOUNDED, GradSampler -from sparseml.pytorch.utils.logger import BaseLogger +from sparseml.pytorch.utils import MEMORY_BOUNDED, BaseLogger, GradSampler from sparseml.utils import interpolate diff --git a/src/sparseml/pytorch/torchvision/train.py b/src/sparseml/pytorch/torchvision/train.py index bb77f6b8a7f..60d7efeee46 100644 --- a/src/sparseml/pytorch/torchvision/train.py +++ b/src/sparseml/pytorch/torchvision/train.py @@ -50,17 +50,17 @@ from sparseml.pytorch.optim import ScheduledModifierManager from sparseml.pytorch.torchvision import presets, transforms, utils from sparseml.pytorch.torchvision.sampler import RASampler -from sparseml.pytorch.utils.helpers import ( - default_device, - download_framework_model_by_recipe_type, - torch_distributed_zero_first, -) -from sparseml.pytorch.utils.logger import ( +from sparseml.pytorch.utils import ( LoggerManager, PythonLogger, TensorBoardLogger, WANDBLogger, ) +from sparseml.pytorch.utils.helpers import ( + default_device, + download_framework_model_by_recipe_type, + torch_distributed_zero_first, +) from sparseml.pytorch.utils.model import load_model, model_to_device from sparsezoo import Model diff --git a/src/sparseml/pytorch/utils/__init__.py b/src/sparseml/pytorch/utils/__init__.py index c99bd5638b1..a1ae444389c 100644 --- a/src/sparseml/pytorch/utils/__init__.py +++ b/src/sparseml/pytorch/utils/__init__.py @@ -18,13 +18,13 @@ # flake8: noqa +from ...core.logger import * from ..base import check_torch_install as _check_torch_install from .benchmarker import * from .distributed import * from .exporter import * from .helpers import * from .log_sparsification_info import * -from .logger import * from .loss import * from .model import * from .module import * diff --git a/src/sparseml/pytorch/utils/log_sparsification_info.py b/src/sparseml/pytorch/utils/log_sparsification_info.py index 5679aec70a9..db9d93bad3f 100644 --- a/src/sparseml/pytorch/utils/log_sparsification_info.py +++ b/src/sparseml/pytorch/utils/log_sparsification_info.py @@ -16,7 +16,7 @@ import torch -from sparseml.pytorch.utils.logger import BaseLogger +from sparseml.pytorch.utils import BaseLogger from sparseml.pytorch.utils.sparsification_info.module_sparsification_info import ( ModuleSparsificationInfo, ) diff --git a/src/sparseml/pytorch/utils/module.py b/src/sparseml/pytorch/utils/module.py index cbc4c5cccc5..a8200b87d72 100644 --- a/src/sparseml/pytorch/utils/module.py +++ b/src/sparseml/pytorch/utils/module.py @@ -31,13 +31,13 @@ from torch.utils.hooks import RemovableHandle from tqdm import auto +from sparseml.pytorch.utils import BaseLogger from sparseml.pytorch.utils.helpers import ( get_optim_learning_rate, tensors_batch_size, tensors_module_forward, tensors_to_device, ) -from sparseml.pytorch.utils.logger import BaseLogger from sparseml.pytorch.utils.loss import DEFAULT_LOSS_KEY, LossWrapper diff --git a/src/sparseml/yolov8/trainers.py b/src/sparseml/yolov8/trainers.py index b3745b95b96..21a85f2839c 100644 --- a/src/sparseml/yolov8/trainers.py +++ b/src/sparseml/yolov8/trainers.py @@ -30,9 +30,13 @@ from sparseml.optim.helpers import load_recipe_yaml_str from sparseml.pytorch.optim.manager import ScheduledModifierManager from sparseml.pytorch.sparsification.quantization import skip_onnx_input_quantize -from sparseml.pytorch.utils import ModuleExporter +from sparseml.pytorch.utils import ( + LoggerManager, + ModuleExporter, + PythonLogger, + WANDBLogger, +) from sparseml.pytorch.utils.helpers import download_framework_model_by_recipe_type -from sparseml.pytorch.utils.logger import LoggerManager, PythonLogger, WANDBLogger from sparseml.yolov8.modules import Bottleneck, Conv from sparseml.yolov8.utils import ( check_coco128_segmentation, diff --git a/tests/sparseml/pytorch/utils/test_logger.py b/tests/sparseml/core/test_logger.py similarity index 98% rename from tests/sparseml/pytorch/utils/test_logger.py rename to tests/sparseml/core/test_logger.py index 7cceeff3017..d26fe3bee42 100644 --- a/tests/sparseml/pytorch/utils/test_logger.py +++ b/tests/sparseml/core/test_logger.py @@ -19,7 +19,7 @@ import pytest -from sparseml.pytorch.utils import ( +from sparseml.core.logger import ( LambdaLogger, LoggerManager, PythonLogger,