Pull request #5: Quantization Helper & pyproject improvement

Merge in CSID/glaucus from feature/toml-and-utils to main

Squashed commit of the following:

commit 631df6e10625bdef6580e0da333d0b43704e34f2
Author: Kyle A Logue <kyle.a.logue@aero.org>
Date:   Thu Feb 8 15:57:03 2024 -0800

    Quantization Helper & pyproject improvement

    * move all configuration into pyproject.toml
    * add function to adapt quantized weights to non-quantized model
    * increment to v1.1.4

README.md

@@ -3,7 +3,8 @@
 # Glaucus
 
 The Aerospace Corporation is proud to present our complex-valued encoder,
-decoder, and a new loss function for RF DSP in PyTorch.
+decoder, and a new loss function for radio frequency (RF) digital signal
+processing (DSP) in PyTorch.
 
 ## Video (click to play)
 
@@ -18,8 +19,9 @@ decoder, and a new loss function for RF DSP in PyTorch.
 
 ### Testing
 
-* `coverage run -a --source=glaucus -m pytest --doctest-modules; coverage html`
-* `pytest .`
+* `pytest`
+* `coverage run`
+* `pylint glaucus tests`
 
 ### Use pre-trained model with SigMF data
 
@@ -41,6 +43,7 @@ state_dict = torch.hub.load_state_dict_from_url(
  map_location='cpu')
 model.load_state_dict(state_dict)
 # prepare for prediction
+model.freeze()
 model.eval()
 torch.quantization.convert(model, inplace=True)
 # get samples into NL tensor
@@ -53,6 +56,7 @@ y_encoded_uint8 = torch.int_repr(y_encoded)
 ```
 
 #### Higher-accuracy pre-trained model
+
 ```python
 # define architecture
 import torch
@@ -71,6 +75,24 @@ model.load_state_dict(state_dict)
 # see above for rest
 ```
 
+#### Use pre-trained model & discard quantization layers
+
+```python
+# create model, but skip quantization
+from glaucus.utils import adapt_glaucus_quantized_weights
+model = GlaucusAE(bottleneck_quantize=False, data_format='nl')
+state_dict = torch.hub.load_state_dict_from_url(
+ 'https://github.com/the-aerospace-corporation/glaucus/releases/download/v1.1.0/glaucus-512-3275-5517642b.pth',
+ map_location='cpu')
+state_dict = adapt_glaucus_quantized_weights(state_dict)
+# ignore "unexpected_keys" warning
+model.load_state_dict(state_dict, strict=False)
+# prepare for evaluation mode
+model.freeze()
+model.eval()
+# see above for rest
+```
+
 ### Get loss between two RF signals
 
 ```python

glaucus/__init__.py

@@ -2,10 +2,11 @@
 # This file is a part of Glaucus
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-__version__ = '1.1.3'
+__version__ = '1.1.4'
 
 from .rfloss import *
 from .layers import *
 from .gblocks import *
 from .fcblocks import *
 from .autoencoders import *
+from .utils import *

glaucus/utils.py

@@ -0,0 +1,44 @@
+'''utilities'''
+# Copyright 2023 The Aerospace Corporation
+# This file is a part of Glaucus
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+import copy
+import re
+
+
+def adapt_glaucus_quantized_weights(state_dict: dict) -> dict:
+ """
+ The pretrained Glaucus models have a quantization layer that shifts the
+ encoder list positions, so if we create a model w/o quantization we have to
+ shift those layers slightly to make the pretrained model work.
+
+ This function decrements the position of the decoder layers in the state
+ dict to allow loading from a pre-trained model that was quantization aware.
+
+ ie: `fc_decoder._fc.1.weight` becomes `fc_decoder._fc.0.weight`
+
+ There will be extra layers remaining, but we can discard them by loading
+ with `strict=False`. See the README for an example.
+
+ Parameters
+ ----------
+ state_dict : dict
+ Torch state dictionary including quantization layers.
+
+ Returns
+ -------
+ new_state_dict : dict
+ State dictionary without quantization layers.
+ """
+ new_state_dict = copy.deepcopy(state_dict)
+
+ pattern = r"(fc_decoder._fc.)(\d+)(\.\w+)" # regex pattern
+
+ for key, value in state_dict.items():
+ match = re.match(pattern, key)
+ if match:
+ extracted_int = int(match.group(2))
+ new_key = f"{match.group(1)}{extracted_int-1}{match.group(3)}"
+ new_state_dict[new_key] = value
+ return new_state_dict

pyproject.toml

@@ -0,0 +1,76 @@
+[project]
+name = "glaucus"
+description = "Glaucus is a PyTorch complex-valued ML autoencoder & RF estimation python module. "
+keywords = ["dsp", "ml", "autoencoder", "sigint", "rf"]
+classifiers = [
+ "License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)",
+ "Operating System :: OS Independent",
+ "Programming Language :: Python :: 3",
+ "Programming Language :: Python :: 3.8",
+ "Programming Language :: Python :: 3.9",
+ "Programming Language :: Python :: 3.10",
+ "Programming Language :: Python :: 3.11",
+ "Programming Language :: Python :: 3.12",
+]
+dynamic = ["version", "readme"]
+authors = [
+ {name = "Kyle Logue", email = "kyle.logue@aero.org"}
+]
+requires-python = ">=3.8"
+dependencies = [
+ "torch", # basic ML framework
+ "lightning", # extensions for PyTorch
+ "madgrad", # our favorite optimizer
+ "hypothesis", # best unit testing
+ ]
+ [project.urls]
+ repository = "https://github.com/the-aerospace-corporation/glaucus"
+
+[tool.setuptools]
+packages = ["glaucus"]
+ [tool.setuptools.dynamic]
+ version = {attr = "glaucus.__version__"}
+ readme = {file = ["README.md"], content-type = "text/markdown"}
+
+[build-system]
+requires = ["setuptools>=65.0", "setuptools-scm"]
+build-backend = "setuptools.build_meta"
+
+[tool.coverage.run]
+branch = true
+source = ["glaucus", "tests"]
+# -rA captures stdout from all tests and places it after the pytest summary
+command_line = "-m pytest -rA --doctest-modules --junitxml=pytest.xml"
+
+[tool.pytest.ini_options]
+addopts = "--doctest-modules"
+testpaths = ["glaucus", "tests"]
+
+[tool.pylint]
+ [tool.pylint.main]
+ load-plugins = [
+ "pylint.extensions.typing",
+ "pylint.extensions.docparams",
+ ]
+ exit-zero = true
+ [tool.pylint.messages_control]
+ disable = [
+ "logging-not-lazy",
+ "missing-module-docstring",
+ "import-error",
+ "unspecified-encoding",
+ ]
+ max-line-length = 160
+ [tool.pylint.REPORTS]
+ # omit from the similarity reports
+ ignore-comments = "yes"
+ ignore-docstrings = "yes"
+ ignore-imports = "yes"
+ ignore-signatures = "yes"
+ min-similarity-lines = 4
+
+[tool.pytype]
+inputs = ["glaucus", "tests"]
+
+[tool.black]
+line-length = 160

tests/test_autoencoders.py

@@ -1,23 +1,23 @@
-'''ensure autoencoders are working'''
+"""ensure autoencoders are working"""
 # Copyright 2023 The Aerospace Corporation
 # This file is a part of Glaucus
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 import unittest
 import torch
 
-from glaucus import GlaucusAE, FullyConnectedAE
+from glaucus import FullyConnectedAE, GlaucusAE
 
 
 class TestAE(unittest.TestCase):
  def test_ae_roundtrip(self):
- '''the output size should always be the same as the input size'''
+ """the output size should always be the same as the input size"""
  for AE in [GlaucusAE, FullyConnectedAE]:
- for data_format in ['ncl', 'nl']:
- for domain in ['time', 'freq']:
+ for data_format in ["ncl", "nl"]:
+ for domain in ["time", "freq"]:
  # note if we use a diff spatial_size, will need to gen new encoder & decoder bocks
  spatial_size = 4096
- if data_format == 'ncl':
+ if data_format == "ncl":
  trash_x = torch.randn(7, 2, spatial_size)
  else:
  trash_x = torch.randn(7, spatial_size, dtype=torch.complex64)
@@ -26,14 +26,14 @@ def test_ae_roundtrip(self):
  self.assertEqual(trash_x.shape, trash_y.shape)
 
  def test_ae_quantization(self):
- '''If quantization enabled, should use quint8 as latent output'''
+ """If quantization enabled, should use quint8 as latent output"""
  for AE in [FullyConnectedAE, GlaucusAE]:
- for data_format in ['ncl', 'nl']:
+ for data_format in ["ncl", "nl"]:
  for is_quantized in [True, False]:
  target = torch.quint8 if is_quantized else torch.float32
  # note if we use a diff spatial_size, will need to gen new encoder & decoder bocks
  spatial_size = 4096
- if data_format == 'ncl':
+ if data_format == "ncl":
  trash_x = torch.randn(7, 2, spatial_size)
  else:
  trash_x = torch.randn(7, spatial_size, dtype=torch.complex64)
@@ -47,13 +47,13 @@ def test_ae_quantization(self):
  self.assertEqual(trash_latent.dtype, target)
 
  def test_ae_backprop(self):
- '''catch errors during backpropagation'''
- for data_format in ['ncl', 'nl']:
+ """catch errors during backpropagation"""
+ for data_format in ["ncl", "nl"]:
  for AE in [FullyConnectedAE, GlaucusAE]:
  for is_quantized in [True, False]:
  # note if we use a diff spatial_size, will need to gen new encoder & decoder bocks
  spatial_size = 4096
- if data_format == 'ncl':
+ if data_format == "ncl":
  trash_x = torch.randn(7, 2, spatial_size)
  else:
  trash_x = torch.randn(7, spatial_size, dtype=torch.complex64)

tests/test_blocks.py

@@ -1,4 +1,4 @@
-'''ensure blocks are working'''
+"""ensure blocks are working"""
 # Copyright 2023 The Aerospace Corporation
 # This file is a part of Glaucus
 # SPDX-License-Identifier: LGPL-3.0-or-later
@@ -7,25 +7,26 @@
 import torch
 from hypothesis import settings, given, strategies as st
 
-from glaucus import GlaucusNet, blockgen, FullyConnected, GBlock
+from glaucus import FullyConnected, GBlock, GlaucusNet, blockgen
 
 
 class TestParams(unittest.TestCase):
- '''autoencoders should operate over all valid params'''
+ """autoencoders should operate over all valid params"""
+
  @settings(deadline=None, max_examples=100)
  @given(
  exponent=st.integers(min_value=2, max_value=14),
  steps=st.integers(min_value=1, max_value=8),
  filters_mid=st.integers(min_value=1, max_value=100)
  )
  def test_io_glaucusnet(self, exponent, steps, filters_mid):
- '''design works on a variety of spatial sizes'''
+ """design works on a variety of spatial sizes"""
  spatial_dim = 2**exponent
  # for spatial_dim in 2**np.arange(8, 14):
- encoder_blocks = blockgen(steps=steps, spatial_in=spatial_dim, spatial_out=8, filters_in=2, filters_out=filters_mid, mode='encoder')
- decoder_blocks = blockgen(steps=steps, spatial_in=8, spatial_out=spatial_dim, filters_in=filters_mid, filters_out=2, mode='decoder')
- encoder = GlaucusNet(mode='encoder', blocks=encoder_blocks, spatial_dim=spatial_dim)
- decoder = GlaucusNet(mode='decoder', blocks=decoder_blocks, spatial_dim=spatial_dim)
+ encoder_blocks = blockgen(steps=steps, spatial_in=spatial_dim, spatial_out=8, filters_in=2, filters_out=filters_mid, mode="encoder")
+ decoder_blocks = blockgen(steps=steps, spatial_in=8, spatial_out=spatial_dim, filters_in=filters_mid, filters_out=2, mode="decoder")
+ encoder = GlaucusNet(mode="encoder", blocks=encoder_blocks, spatial_dim=spatial_dim)
+ decoder = GlaucusNet(mode="decoder", blocks=decoder_blocks, spatial_dim=spatial_dim)
  trash_x = torch.randn(3, 2, spatial_dim)
  trash_y = decoder(encoder(trash_x))
  self.assertEqual(trash_x.shape, trash_y.shape)
@@ -53,9 +54,7 @@ def test_io_gblock(self, spatial_exponent, filters_in, filters_out, stride, expa
  squeeze_ratio -= 1
  squeeze_ratio = max(filters_in * expand_ratio, squeeze_ratio)
  blk = GBlock(
- filters_in=filters_in, filters_out=filters_out,
- stride=stride, kernel_size=kernel_size,
- expand_ratio=expand_ratio, squeeze_ratio=squeeze_ratio
+ filters_in=filters_in, filters_out=filters_out, stride=stride, kernel_size=kernel_size, expand_ratio=expand_ratio, squeeze_ratio=squeeze_ratio
  )
  trash_x = torch.randn(2, filters_in, spatial_size)
  trash_y = blk(trash_x)
@@ -71,16 +70,14 @@ def test_io_gblock(self, spatial_exponent, filters_in, filters_out, stride, expa
  exponent_in=st.integers(min_value=2, max_value=14),
  exponent_out=st.integers(min_value=2, max_value=14),
  steps=st.integers(min_value=1, max_value=5),
- quantize_in=st.booleans(), quantize_out=st.booleans(),
- use_dropout=st.booleans()
+ quantize_in=st.booleans(),
+ quantize_out=st.booleans(),
+ use_dropout=st.booleans(),
  )
  def test_io_fc(self, exponent_in, exponent_out, steps, quantize_in, quantize_out, use_dropout):
- '''block should work with a variety of configs'''
+ """block should work with a variety of configs"""
  size_in, size_out = exponent_in**2, exponent_out**2
- autoencoder = FullyConnected(
- size_in=size_in, size_out=size_out,
- steps=steps, quantize_in=quantize_in, quantize_out=quantize_out
- )
+ autoencoder = FullyConnected(size_in=size_in, size_out=size_out, steps=steps, quantize_in=quantize_in, quantize_out=quantize_out)
  trash_x = torch.randn(3, size_in)
  trash_y = autoencoder(trash_x)
  self.assertEqual(trash_y.shape[1], size_out)