Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/add_randomSolution'
Browse files Browse the repository at this point in the history
  • Loading branch information
amkrajewski committed Feb 13, 2024
2 parents 56450aa + e2c9e92 commit 23c73f7
Show file tree
Hide file tree
Showing 17 changed files with 1,510 additions and 696 deletions.
3 changes: 1 addition & 2 deletions .github/workflows/benchmarks.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,8 @@ jobs:
strategy:
matrix:
python-version: ["3.9", "3.10", "3.11"]
platform: [ubuntu-latest]
fail-fast: false
runs-on: ${{ matrix.platform }}
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@v3

Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/coreTests_LinuxUbuntu.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ jobs:
run: |
python -m pip install --upgrade pip
python -m pip install --upgrade setuptools
python -m pip install wheel flask pytest pytest-cov
python -m pip install wheel flask pytest
python -m pip install -e .
- name: Test with pytest
run: |
Expand Down
16 changes: 8 additions & 8 deletions .github/workflows/fullTest.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ on:

jobs:
testPython309:
runs-on: macos-latest
runs-on: macos-12
strategy:
fail-fast: false
steps:
Expand All @@ -35,7 +35,7 @@ jobs:
run: python -c "import pysipfenn; c = pysipfenn.Calculator(); c.downloadModels(); c.loadModels();"

- name: Test with pytest
run: pytest --cov=pysipfenn --cov-report=xml
run: pytest --durations=0 --cov=pysipfenn --cov-report=xml
env:
MODELS_FETCHED: true

Expand All @@ -44,7 +44,7 @@ jobs:

testPython310:
needs: testPython309
runs-on: macos-latest
runs-on: macos-12
steps:
- uses: actions/checkout@v3

Expand All @@ -59,19 +59,19 @@ jobs:
run: |
python -m pip install --upgrade pip
python -m pip install --upgrade setuptools
python -m pip install wheel flask pytest pytest-cov
python -m pip install wheel flask pytest
python -m pip install -e ".[dev]"
- name: Download Models
run: python -c "import pysipfenn; c = pysipfenn.Calculator(); c.downloadModels(); c.loadModels();"

- name: Test with pytest
run: pytest --cov=pysipfenn --cov-report=xml
run: pytest --durations=0
env:
MODELS_FETCHED: true

testPython311:
needs: testPython310
runs-on: macos-latest
runs-on: macos-12
steps:
- uses: actions/checkout@v3

Expand All @@ -86,12 +86,12 @@ jobs:
run: |
python -m pip install --upgrade pip
python -m pip install --upgrade setuptools
python -m pip install wheel flask pytest pytest-cov
python -m pip install wheel flask pytest
python -m pip install -e ".[dev]"
- name: Download Models
run: python -c "import pysipfenn; c = pysipfenn.Calculator(); c.downloadModels(); c.loadModels();"

- name: Test with pytest
run: pytest --cov=pysipfenn --cov-report=xml
run: pytest --durations=0
env:
MODELS_FETCHED: true
2 changes: 1 addition & 1 deletion .github/workflows/publishPyPI.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ on:

jobs:
deploy:
runs-on: ubuntu-latest
runs-on: macos-14
steps:
- uses: actions/checkout@v3

Expand Down
6 changes: 3 additions & 3 deletions .github/workflows/weeklyTesting.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,11 @@ jobs:
python -m pip install -e .
- name: Test with pytest
run: |
pytest
pytest --durations=0
testFullPython310:
needs: coreTests
runs-on: macos-latest
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@v3

Expand All @@ -55,7 +55,7 @@ jobs:
run: python -c "import pysipfenn; c = pysipfenn.Calculator(); c.downloadModels(); c.loadModels();"

- name: Test with pytest
run: pytest
run: pytest --durations=0
env:
MODELS_FETCHED: true

Expand Down
13 changes: 7 additions & 6 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -28,22 +28,23 @@ readme = "README.md"

requires-python = ">=3.9"
dependencies = [
"pymatgen>=2023.2.22",
"torch>=1.11.0",
"pymatgen>=2023.10.3",
"torch>=2.0",
"onnx2torch>=1.5.2",
"onnx>=1.13.0",
"numpy>=1.22.0",
"numpy>=1.25.0",
"tqdm>=4.65.0",
"natsort>=8.3.0",
"pymongo>=4.2",
"pymongo>=4.4",
"pySmartDL>=1.3.4",
"dnspython",
"ruamel.yaml"
]

[project.optional-dependencies]
dev = [
"coremltools>=6.3",
"onnxconverter_common>=1.13.0",
"coremltools>=7.0",
"onnxconverter_common>=1.14.0",
"onnxsim==0.4.33",
"onnxruntime>=1.16.0"
]
Expand Down
92 changes: 53 additions & 39 deletions pysipfenn/core/modelExporters.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,22 +15,22 @@

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.
(2) simplify the models using ONNX optimizer, and (3) convert them to `FP16` precision, cutting the size in half.
Args:
calculator: A calculator object with loaded models that has loaded PyTorch models (happens automatically
when the autoLoad argument is kept to its default value of True when initializing the Calculator). During the
calculator: A ``Calculator`` object with loaded models that has loaded PyTorch models (happens automatically
when the ``autoLoad`` argument is kept to its default value of ``True`` when initializing the Calculator). During the
initialization, the loaded PyTorch models are converted back to ONNX (in memory) to be then either adjusted or
persisted to disk.
Attributes:
calculator: A calculator object with ONNX loaded models.
simplifiedDict: A dictionary of models that have been simplified.
fp16Dict: A dictionary of models that have been converted to FP16.
calculator: A Calculator object with ONNX loaded models.
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."""
"""Initialize the ``ONNXExporter`` using a calculator object."""
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 @@ -80,7 +80,7 @@ def simplify(self, model: str) -> None:
"""Simplify a loaded model using the ONNX optimizer.
Args:
model: The name of the model to simplify (must be loaded in the Calculator).
model: The name of the model to simplify (must be loaded in the ``Calculator``).
Returns:
None
Expand All @@ -100,11 +100,11 @@ def simplifyAll(self):
self.simplify(model)
print('***** Done simplifying all models! *****')

def toFP16(self, model: str):
def toFP16(self, model: str) -> None:
"""Convert a loaded model to FP16 precision.
Args:
model: The name of the model to convert to FP16 (must be loaded in the Calculator).
model: The name of the model to convert to FP16 (must be loaded in the ``Calculator``).
Returns:
None
Expand All @@ -124,11 +124,13 @@ def toFP16All(self):
self.toFP16(model)
print('***** Done converting all models to FP16! *****')

def export(self, model: str):
"""Export a loaded model to ONNX format.
def export(self, model: str, append: str = '') -> None:
"""Export a loaded model to ``ONNX``format.
Args:
model: The name of the model to export (must be loaded in the Calculator).
model: The name of the model to export (must be loaded in the ``Calculator``).
append: A string to append to the exported model name after the model name, simplification marker, and
FP16 marker. Useful for adding a version number or other information to the exported model name.
Returns:
None
Expand All @@ -141,40 +143,46 @@ def export(self, model: str):
name += '_simplified'
if self.fp16Dict[model]:
name += '_fp16'
if append:
name += f'_{append}'
name += '.onnx'
onnx.save(loadedModel, name)
print(f'--> Exported as {name}', flush=True)

def exportAll(self):
"""Export all loaded models to ONNX format with the export function."""
def exportAll(self, append: str = '') -> None:
"""Export all loaded models to ``ONNX`` format with the export function. ``append`` string can be passed to the export
function to append to the exported model name.
"""
for model in tqdm(self.calculator.loadedModels):
self.export(model)
self.export(model, append=append)
print('***** Done exporting all models! *****')


class TorchExporter:
"""Export models to the PyTorch PT format to allow for easy loading and inference in PyTorch in other projects.
"""Export models to the ``PyTorch PT`` format to allow for easy loading and inference in PyTorch in other projects.
Args:
calculator: A calculator object with loaded models.
calculator: A ``Calculator`` object with loaded models.
Attributes:
calculator: A calculator object with loaded models.
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
assert len(self.calculator.loadedModels) > 0, 'No models loaded in calculator. Nothing to export.'
print(f'Initialized TorchExporter with models: {list(self.calculator.loadedModels.keys())}')

def export(self, model: str):
"""Export a loaded model to PyTorch PT format. Models are exported in eval mode (no dropout) and saved in the
def export(self, model: str, append: str = '') -> None:
"""Export a loaded model to ``PyTorch PT`` format. Models are exported in eval mode (no dropout) and saved in the
current working directory.
Args:
model: The name of the model to export (must be loaded in the Calculator) and it must have a descriptor
(Ward2017 or KS2022) defined in the calculator.models dictionary created when the Calculator was
model: The name of the model to export (must be loaded in the ``Calculator``) and it must have a descriptor
(``Ward2017`` or ``KS2022``) defined in the ``Calculator.models`` dictionary created when the ``Calculator`` was
initialized.
append: A string to append to the exported model name after the model name. Useful for adding a version
number or other information to the exported model name.
Returns:
None
Expand All @@ -200,43 +208,47 @@ def export(self, model: str):

tracedModel = torch.jit.trace(loadedModel, inputs_tracer)

name = f"{model}.pt"
name = f"{model}{f'_{append}' if append else ''}.pt"
tracedModel.save(name)
print(f'--> Exported as {name}', flush=True)

def exportAll(self):
"""Export all loaded models to PyTorch PT format with the export function."""
def exportAll(self, append: str = '') -> None:
"""Exports all loaded models to PyTorch PT format with the export function. `append` can be passed to the export
function
"""
for model in tqdm(self.calculator.loadedModels):
self.export(model)
self.export(model, append=append)
print('***** Done exporting all models! *****')


class CoreMLExporter:
"""Export models to the CoreML format to allow for easy loading and inference in CoreML in other projects,
"""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.
Args:
calculator: A calculator object with loaded models.
calculator: A ``Calculator`` object with loaded models.
Attributes:
calculator: A calculator object with loaded models.
calculator: A ``Calculator`` object with loaded models.
"""
def __init__(self, calculator: Calculator):
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())}')

def export(self, model: str):
"""Export a loaded model to CoreML format. Models will be saved as {model}.mlpackage in the current working
directory. Models will be annotated with the feature vector name (Ward2017 or KS2022) and the output will be
def export(self, model: str, append: str = '') -> None:
"""Export a loaded model to ``CoreML`` format. Models will be saved as ``{model}.mlpackage`` in the current working
directory. Models will be annotated with the feature vector name (``Ward2017`` or ``KS2022``) and the output will be
named "property". The latter behavior will be adjusted in the future when model output name and unit will be
added to the model JSON metadata.
Args:
model: The name of the model to export (must be loaded in the Calculator) and it must have a descriptor
(Ward2017 or KS2022) defined in the calculator.models dictionary created when the Calculator was
model: The name of the model to export (must be loaded in the ``Calculator``) and it must have a descriptor
(``Ward2017`` or ``KS2022``) defined in the ``calculator.models`` dictionary created when the ``Calculator`` was
initialized.
append: A string to append to the exported model name after the model name. Useful for adding a version
number or other information to the exported model name.
Returns:
None
Expand Down Expand Up @@ -270,12 +282,14 @@ def export(self, model: str):
inputs=inputs_converter,
outputs=[ct.TensorType(name='property')]
)
name = f"{model}.mlpackage"
name = f"{model}{f'_{append}' if append else ''}.mlpackage"
coreml_model.save(name)
print(f'--> Exported as {name}', flush=True)

def exportAll(self):
"""Export all loaded models to CoreML format with the export function."""
def exportAll(self, append: str = '') -> None:
"""Export all loaded models to ``CoreML`` format with the export function. ``append`` can be passed to the export
function to append to all exported model names.
"""
for model in tqdm(self.calculator.loadedModels):
self.export(model)
self.export(model, append=append)
print('***** Done exporting all models! *****')
Loading

0 comments on commit 23c73f7

Please sign in to comment.