Skip to content

Commit

Permalink
(ME) distributed class-specific imports to limit displaying of the ex…
Browse files Browse the repository at this point in the history
…tra-dependency requirement to only when needed
  • Loading branch information
amkrajewski committed Feb 17, 2024
1 parent b36330a commit fcd09ca
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 11 deletions.
33 changes: 23 additions & 10 deletions pysipfenn/core/modelExporters.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,18 +4,12 @@
import io
from tqdm import tqdm

try:
import coremltools as ct
from onnxconverter_common import float16
from onnxsim import simplify
except ModuleNotFoundError as e:
print('Note: Export Dependencies are not installed by default. If you need them, you have to install pySIPFENN in '
'"dev" mode like: pip install -e "pysipfenn[dev]", or like pip install -e ".[dev]" (see pysipfenn.org)')


class ONNXExporter:
"""Export models to the ONNX format (what they ship in by default) to allow (1) exporting modified pySIPFENN models,
(2) simplify the models using ONNX optimizer, and (3) convert them to `FP16` precision, cutting the size in half.
Note: Some of the dependencies (``onnxconverter_common`` and ``onnxsim``) are not installed by default. If you need them,
you have to install pySIPFENN in `dev` mode like: ``pip install "pysipfenn[dev]"``, or like ``pip install -e ".[dev]"``.
Args:
calculator: A ``Calculator`` object with loaded models that has loaded PyTorch models (happens automatically
Expand All @@ -28,9 +22,16 @@ class ONNXExporter:
simplifiedDict: A boolean dictionary of models that have been simplified.
fp16Dict: A boolean dictionary of models that have been converted to FP16.
"""

def __init__(self, calculator: Calculator):
"""Initialize the ``ONNXExporter`` using a calculator object."""
try:
from onnxconverter_common import float16
from onnxsim import simplify
except ModuleNotFoundError as e:
raise Exception(str(e) + '\n\nNote: Export Dependencies are not installed by default. If you need them, you have to install '
'pySIPFENN in `dev` mode like: `pip install "pysipfenn[dev]"`, or like `pip install -e ".[dev]"` (see pysipfenn.org)')

self.simplifiedDict = {model: False for model in calculator.loadedModels.keys()}
self.fp16Dict = {model: False for model in calculator.loadedModels.keys()}
self.calculator = calculator
Expand Down Expand Up @@ -167,6 +168,7 @@ class TorchExporter:
Attributes:
calculator: A ``Calculator`` object with loaded models.
"""

def __init__(self, calculator: Calculator):
"""Initialize the TorchExporter with a calculator object that has loaded models."""
self.calculator = calculator
Expand Down Expand Up @@ -225,14 +227,25 @@ class CoreMLExporter:
"""Export models to the ``CoreML`` format to allow for easy loading and inference in ``CoreML`` in other projects,
particularly valuable for Apple devices, as pySIPFENN models can be run using the Neural Engine accelerator
with minimal power consumption and neat optimizations.
Note: Some of the dependencies (``coremltools``) are not installed by default. If you need them,
you have to install pySIPFENN in `dev` mode like: ``pip install "pysipfenn[dev]"``, or like ``pip install -e ".[dev]"``.
Args:
calculator: A ``Calculator`` object with loaded models.
Attributes:
calculator: A ``Calculator`` object with loaded models.
"""

def __init__(self, calculator: Calculator):
try:
import coremltools as ct
except ModuleNotFoundError as e:
raise Exception(str(e) + '\n\nNote: Export Dependencies are not installed by default. If you need them, you have to install '
'pySIPFENN in `dev` mode like: `pip install "pysipfenn[dev]"`, or like `pip install -e ".[dev]"` (see pysipfenn.org)')


self.calculator = calculator
assert len(self.calculator.loadedModels)>0, 'No models loaded in calculator. Nothing to export.'
print(f'Initialized CoreMLExporter with models: {list(self.calculator.loadedModels.keys())}')
Expand Down
2 changes: 1 addition & 1 deletion pysipfenn/descriptorDefinitions/KS2022_randomSolutions.py
Original file line number Diff line number Diff line change
Expand Up @@ -410,7 +410,7 @@ def generate_descriptor(struct: Structure,
diffArray = np.array(diffHistory)

# Plot the parameters as lines. Add hover text to show the parameter name based on the labels_KS2022.csv file.
with resources.files('pysipfenn').joinpath('descriptorDefinitions/labels_KS2022.csv').open() as f:
with resources.files('pysipfenn.descriptorDefinitions').joinpath('labels_KS2022.csv').open() as f:
labels = f.readlines()
fig = px.line(pd.DataFrame(diffArray, columns=labels), title='KS2022 Descriptor Parameters',
range_y=[-0.5, 0.5])
Expand Down

0 comments on commit fcd09ca

Please sign in to comment.