spikingjelly 2.0.0.dev0__tar.gz → 2.0.0.dev1__tar.gz

{spikingjelly-2.0.0.dev0 → spikingjelly-2.0.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spikingjelly
-Version: 2.0.0.dev0
+Version: 2.0.0.dev1
 Summary: A deep learning framework for SNNs built on PyTorch.
 Author: PKU MLG, PCL, and other contributors
 Author-email: fwei@pku.edu.cn, chyq@pku.edu.cn, yfhuang@pku.edu.cn
@@ -20,12 +20,11 @@ License-File: LICENSES/translations/LICENSE.hi
 Requires-Dist: torch>=2.6.0
 Requires-Dist: torchvision
 Requires-Dist: torchaudio
-Requires-Dist: einops
 Requires-Dist: h5py
+Requires-Dist: loguru<1,>=0.7
 Requires-Dist: matplotlib
 Requires-Dist: numpy
 Requires-Dist: packaging
-Requires-Dist: pydantic
 Requires-Dist: requests
 Requires-Dist: SciencePlots
 Requires-Dist: scipy
@@ -43,6 +42,12 @@ Requires-Dist: nirtorch; extra == "nir"
 Provides-Extra: lightning
 Requires-Dist: lightning; extra == "lightning"
 Requires-Dist: jsonargparse[signatures]; extra == "lightning"
+Provides-Extra: fp8-torchao
+Requires-Dist: torchao<0.18,>=0.13; extra == "fp8-torchao"
+Provides-Extra: fp8-te
+Requires-Dist: transformer-engine[pytorch]<3,>=2.16; extra == "fp8-te"
+Provides-Extra: qwen
+Requires-Dist: transformers==5.13.0; extra == "qwen"
 Dynamic: license-file
 
 # SpikingJelly
@@ -208,6 +213,8 @@ generations, `MINOR` adds backward-compatible functionality, and `PATCH` fixes
 bugs. Python package pre-release spelling is used for V2 development releases,
 for example `2.0.0.dev0`, `2.0.0a1`, `2.0.0b1`, and `2.0.0rc1`.
 
+See [CHANGELOG.md](./CHANGELOG.md) for the V2 release changelog.
+
 <details>
 <summary>Compatibility, migration, and older docs</summary>
 

{spikingjelly-2.0.0.dev0 → spikingjelly-2.0.0.dev1}/README.md

@@ -161,6 +161,8 @@ generations, `MINOR` adds backward-compatible functionality, and `PATCH` fixes
 bugs. Python package pre-release spelling is used for V2 development releases,
 for example `2.0.0.dev0`, `2.0.0a1`, `2.0.0b1`, and `2.0.0rc1`.
 
+See [CHANGELOG.md](./CHANGELOG.md) for the V2 release changelog.
+
 <details>
 <summary>Compatibility, migration, and older docs</summary>
 

{spikingjelly-2.0.0.dev0 → spikingjelly-2.0.0.dev1}/pyproject.toml

@@ -10,7 +10,7 @@ include = ["spikingjelly*"]
 
 [project]
 name = "spikingjelly"
-version = "2.0.0.dev0"
+version = "2.0.0.dev1"
 description = "A deep learning framework for SNNs built on PyTorch."
 readme = {file = "README.md", content-type = "text/markdown"}
 authors = [
@@ -24,12 +24,11 @@ dependencies = [
     "torch>=2.6.0",
     "torchvision",
     "torchaudio",
-    "einops",
     "h5py",
+    "loguru>=0.7,<1",
     "matplotlib",
     "numpy",
     "packaging",
-    "pydantic",
     "requests",
     "SciencePlots",
     "scipy",
@@ -50,13 +49,27 @@ cupy12 = ["cupy-cuda12x"]
 triton = ["triton>=3.3.1"]
 nir = ["nir", "nirtorch"]
 lightning = ["lightning", "jsonargparse[signatures]"]
+fp8-torchao = [
+    "torchao>=0.13,<0.18",
+]
+fp8-te = [
+    "transformer-engine[pytorch]>=2.16,<3",
+]
+qwen = [
+    "transformers==5.13.0",
+]
 
 [dependency-groups]
 dev = [
     "pytest",
+    "pytest-cov",
     "pydata-sphinx-theme>=0.16.1",
     "sphinx>=8.0.0,<9",
 ]
+llm-benchmark = [
+    "datasets>=5,<6",
+    "lm-eval[hf]==0.4.12",
+]
 
 [project.urls]
 Homepage = "https://github.com/fangwei123456/spikingjelly"
@@ -73,3 +86,20 @@ quote-style = "double"
 
 [tool.ruff.lint]
 ignore = ["E402", "E731", "E741", "F403"]
+
+[tool.ruff.lint.isort]
+known-first-party = ["benchmark", "spikingjelly", "test"]
+
+[tool.pytest.ini_options]
+pythonpath = ["."]
+testpaths = ["test"]
+
+[tool.coverage.run]
+branch = true
+omit = [
+    "*/examples/*",
+]
+source = ["spikingjelly"]
+
+[tool.coverage.report]
+show_missing = true

spikingjelly-2.0.0.dev1/spikingjelly/__init__.py

@@ -0,0 +1,7 @@
+"""SpikingJelly package.
+
+Import the package logger explicitly from :mod:`spikingjelly.logger` when it is
+needed; it is not re-exported from this package namespace.
+"""
+
+__all__ = []

{spikingjelly-2.0.0.dev0 → spikingjelly-2.0.0.dev1}/spikingjelly/activation_based/ann2snn/__init__.py

@@ -9,10 +9,10 @@ r"""
 
 ANN 到 SNN 的转换模块。提供 FX graph 路径的 :class:`FXConverter` /
 :class:`FXConversionRecipe`、module tree 路径的 :class:`ModuleConverter` /
-:class:`ModuleConversionRecipe`，以及 ``HookFactory``、``NeuronFactory``、
-``ReLURule``、``ThresholdOptimizer`` 等可扩展组件，并附带 ``download_url``
-工具函数。兼容名 :class:`Converter` 等价于 :class:`FXConverter`，
-:class:`ConversionRecipe` 等价于 :class:`FXConversionRecipe`。
+:class:`ModuleConversionRecipe`，以及 ``download_url`` 工具函数。兼容名
+:class:`Converter` 等价于
+:class:`FXConverter`，:class:`ConversionRecipe` 等价于
+:class:`FXConversionRecipe`。
 
 ----
 
@@ -22,30 +22,31 @@ ANN 到 SNN 的转换模块。提供 FX graph 路径的 :class:`FXConverter` /
 
 ANN-to-SNN conversion module. Provides the FX graph :class:`FXConverter` /
 :class:`FXConversionRecipe`, the module-tree :class:`ModuleConverter` /
-:class:`ModuleConversionRecipe`, extensible building blocks —
-:class:`HookFactory`, :class:`NeuronFactory`, :class:`ReLURule` and
-:class:`ThresholdOptimizer` — and a ``download_url`` helper for fetching
-pretrained models. The compatibility name :class:`Converter` is equivalent to
-:class:`FXConverter`, and :class:`ConversionRecipe` is equivalent to
-:class:`FXConversionRecipe`.
+:class:`ModuleConversionRecipe`, and a ``download_url`` helper for fetching
+pretrained models. The compatibility name
+:class:`Converter` is equivalent to :class:`FXConverter`, and
+:class:`ConversionRecipe` is equivalent to :class:`FXConversionRecipe`.
 """
 
 from .converter import Converter, FXConverter, ModuleConverter
 from .delay import estimate_delay_start
-from .factories import HookFactory, NeuronFactory
 from .modules import ChannelVoltageScaler
+from .qcfs import SignedQCFSSequenceEncoder
 from .recipes import (
     ConversionRecipe,
     FXConversionRecipe,
     LocalThresholdBalancingRecipe,
     ModuleConversionRecipe,
+    Qwen2SNNCalibration,
+    Qwen2SNNConfig,
+    Qwen2SNNModel,
+    Qwen2SNNRecipe,
     RateCodingRecipe,
     SpikeZIPTFQANNRecipe,
     STATransformerRecipe,
     TransformerTDEquivalentRecipe,
+    calibrate_qwen2_snn,
 )
-from .rules import ReLURule
-from .threshold import ThresholdOptimizer
 from .utils import download_url
 
 __all__ = [
@@ -55,16 +56,18 @@ __all__ = [
     "ConversionRecipe",
     "FXConversionRecipe",
     "ModuleConversionRecipe",
+    "Qwen2SNNCalibration",
+    "Qwen2SNNConfig",
+    "Qwen2SNNModel",
+    "Qwen2SNNRecipe",
+    "calibrate_qwen2_snn",
     "RateCodingRecipe",
     "LocalThresholdBalancingRecipe",
     "SpikeZIPTFQANNRecipe",
     "STATransformerRecipe",
     "TransformerTDEquivalentRecipe",
     "ChannelVoltageScaler",
+    "SignedQCFSSequenceEncoder",
     "estimate_delay_start",
     "download_url",
-    "ReLURule",
-    "NeuronFactory",
-    "HookFactory",
-    "ThresholdOptimizer",
 ]

{spikingjelly-2.0.0.dev0 → spikingjelly-2.0.0.dev1}/spikingjelly/activation_based/ann2snn/converter.py

@@ -1,4 +1,6 @@
-import warnings
+import threading
+import time
+import types
 from typing import Optional, Union
 
 import torch
@@ -10,6 +12,45 @@ from spikingjelly.activation_based.ann2snn.recipes import (
     ModuleConversionRecipe,
     TransformerTDEquivalentRecipe,
 )
+from spikingjelly.logger import logger
+
+
+_FX_TRACE_LOCK = threading.RLock()
+
+
+def _resolve_device(
+    ann: nn.Module, configured_device: Optional[Union[torch.device, str]]
+) -> torch.device:
+    if configured_device is not None:
+        return torch.device(configured_device)
+    parameter = next(ann.parameters(), None)
+    return parameter.device if parameter is not None else torch.device("cpu")
+
+
+def _symbolic_trace(root: nn.Module) -> fx.GraphModule:
+    with _FX_TRACE_LOCK:
+        original_reshape = torch.reshape
+
+        def proxy_aware_reshape(input, shape):
+            if isinstance(input, fx.Proxy):
+                return input.tracer.create_proxy(
+                    "call_function", original_reshape, (input, shape), {}
+                )
+            return original_reshape(input, shape)
+
+        # Torch 2.6/2.7 validates reshape's shape before dispatching FX Proxy inputs.
+        torch.reshape = proxy_aware_reshape
+        try:
+            tracer = fx.Tracer()
+            graph = tracer.trace(root)
+        finally:
+            torch.reshape = original_reshape
+
+    # Torch 2.6/2.7 FX codegen cannot register repeated PEP 604 union types.
+    for node in graph.nodes:
+        if isinstance(node.type, types.UnionType):
+            node.type = None
+    return fx.GraphModule(tracer.root, graph)
 
 
 class FXConverter:
@@ -79,14 +120,6 @@ class FXConverter:
                 "FXConverter/Converter requires an FXConversionRecipe. "
                 "Use ModuleConverter for ModuleConversionRecipe instances."
             )
-        if recipe == "transformer_spike_equivalent":
-            warnings.warn(
-                "The 'transformer_spike_equivalent' recipe string is deprecated; "
-                "use 'transformer_td_equivalent' instead.",
-                DeprecationWarning,
-                stacklevel=3,
-            )
-            return TransformerTDEquivalentRecipe()
         if recipe == "transformer_td_equivalent":
             return TransformerTDEquivalentRecipe()
         if recipe == "rate_coding":
@@ -107,14 +140,6 @@ class FXConverter:
             f"instance, but got {type(recipe).__name__}."
         )
 
-    def _resolve_device(self, ann: nn.Module) -> torch.device:
-        if self.device is not None:
-            return torch.device(self.device)
-        try:
-            return next(ann.parameters()).device
-        except StopIteration:
-            return torch.device("cpu")
-
     def convert(self, ann: nn.Module) -> nn.Module:
         r"""
         **API Language** - :ref:`中文 <Converter.convert-cn>` | :ref:`English <Converter.convert-en>`
@@ -154,20 +179,31 @@ class FXConverter:
         """
         configured_device = self.device
         original_training_modes: dict[nn.Module, bool] = {}
+        start_time = time.perf_counter()
+        target_device = None
         try:
             original_training_modes = {
                 module: module.training for module in ann.modules()
             }
-            self.device = self._resolve_device(ann)
+            self.device = _resolve_device(ann, self.device)
+            target_device = self.device
             with torch.no_grad():
                 self.recipe.validate(self)
                 ann = self.recipe.before_trace(self, ann)
-                fx_model = fx.symbolic_trace(ann).to(self.device)
+                fx_model = _symbolic_trace(ann).to(self.device)
                 fx_model = self.recipe.after_trace(self, fx_model)
                 fx_model = self.recipe.insert_observers(self, fx_model)
                 fx_model = self.recipe.calibrate(self, fx_model)
                 fx_model = self.recipe.replace(self, fx_model)
                 fx_model = self.recipe.finalize(self, fx_model)
+                logger.info(
+                    "Conversion completed: converter={} recipe={} device={} modules={} elapsed_ms={:.3f}",
+                    type(self).__name__,
+                    type(self.recipe).__name__,
+                    target_device,
+                    sum(1 for _ in fx_model.modules()),
+                    (time.perf_counter() - start_time) * 1000.0,
+                )
                 return fx_model
         finally:
             for module, training in original_training_modes.items():
@@ -239,14 +275,6 @@ class ModuleConverter:
         self.recipe = recipe
         self.device = device
 
-    def _resolve_device(self, ann: nn.Module) -> torch.device:
-        if self.device is not None:
-            return torch.device(self.device)
-        try:
-            return next(ann.parameters()).device
-        except StopIteration:
-            return torch.device("cpu")
-
     def convert(self, ann: nn.Module) -> nn.Module:
         r"""
         **API Language** - :ref:`中文 <ModuleConverter.convert-cn>` | :ref:`English <ModuleConverter.convert-en>`
@@ -279,11 +307,14 @@ class ModuleConverter:
         """
         configured_device = self.device
         original_training_modes: dict[nn.Module, bool] = {}
+        start_time = time.perf_counter()
+        target_device = None
         try:
             original_training_modes = {
                 module: module.training for module in ann.modules()
             }
-            self.device = self._resolve_device(ann)
+            self.device = _resolve_device(ann, self.device)
+            target_device = self.device
             with torch.no_grad():
                 self.recipe.validate(self)
                 converted = self.recipe.convert_module(self, ann)
@@ -293,7 +324,16 @@ class ModuleConverter:
                         "a torch.nn.Module, got "
                         f"{type(converted).__name__}."
                     )
-                return converted.to(self.device)
+                converted = converted.to(self.device)
+                logger.info(
+                    "Conversion completed: converter={} recipe={} device={} modules={} elapsed_ms={:.3f}",
+                    type(self).__name__,
+                    type(self.recipe).__name__,
+                    target_device,
+                    sum(1 for _ in converted.modules()),
+                    (time.perf_counter() - start_time) * 1000.0,
+                )
+                return converted
         finally:
             for module, training in original_training_modes.items():
                 module.training = training

{spikingjelly-2.0.0.dev0 → spikingjelly-2.0.0.dev1}/spikingjelly/activation_based/ann2snn/delay.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import math
+import time
 import warnings
-from typing import Dict, Iterable, List, Optional, Tuple, Union
+from typing import Dict, Iterable, List, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -12,19 +13,15 @@ from spikingjelly.activation_based.ann2snn.modules import (
     ChannelVoltageScaler,
     VoltageScaler,
 )
+from spikingjelly.activation_based.functional.net_config import reset_net
 from spikingjelly.activation_based.neuron.base_node import BaseNode
+from spikingjelly.logger import logger
 
 
 Scaler = Union[VoltageScaler, ChannelVoltageScaler]
 _MIN_READOUT_STEPS = 4
 
 
-def _reset_snn(model: nn.Module) -> None:
-    for module in model.modules():
-        if hasattr(module, "reset"):
-            module.reset()
-
-
 def _extract_batch_input(batch):
     if isinstance(batch, torch.Tensor):
         return batch
@@ -42,18 +39,6 @@ def _as_runtime_tensor(value, x: torch.Tensor) -> torch.Tensor:
     return torch.as_tensor(value, device=x.device, dtype=x.dtype)
 
 
-def _scale_view(scaler: Scaler, x: torch.Tensor) -> torch.Tensor:
-    if isinstance(scaler, ChannelVoltageScaler):
-        return scaler._view_scale(x).to(device=x.device, dtype=x.dtype)
-    return scaler.scale.to(device=x.device, dtype=x.dtype)
-
-
-def _channel_dim(scaler: Scaler) -> Optional[int]:
-    if isinstance(scaler, ChannelVoltageScaler):
-        return scaler.channel_dim
-    return None
-
-
 def _compute_delay_ratio(
     module: BaseNode,
     post_scaler: Scaler,
@@ -65,21 +50,22 @@ def _compute_delay_ratio(
     if not torch.isfinite(v_threshold).all() or (v_threshold <= 0).any():
         raise ValueError("Delay estimation requires finite positive v_threshold.")
 
-    try:
-        v_init = module.get_reset_value("v")
-    except (KeyError, AttributeError):
-        v_init = getattr(module, "v_reset", 0.0)
+    v_init = module.get_reset_value("v")
     if v_init is None:
         v_init = 0.0
     v_init = _as_runtime_tensor(v_init, x)
 
-    scale = _scale_view(post_scaler, x)
+    if isinstance(post_scaler, ChannelVoltageScaler):
+        scale = post_scaler._view_scale(x).to(device=x.device, dtype=x.dtype)
+        channel_dim = post_scaler.channel_dim
+    else:
+        scale = post_scaler.scale.to(device=x.device, dtype=x.dtype)
+        channel_dim = None
     x_nonnegative = torch.clamp(x.detach(), min=0)
     original_activation = x_nonnegative * scale / v_threshold
     required_charge = torch.clamp(v_threshold - v_init, min=0)
     original_required_charge = required_charge * scale / v_threshold
 
-    channel_dim = _channel_dim(post_scaler)
     if channel_dim is None or scale.dim() == 0:
         max_mean_activation = original_activation.mean()
         reduced_required_charge = original_required_charge.mean()
@@ -210,8 +196,13 @@ def estimate_delay_start(
     if num_batches <= 0:
         raise ValueError("num_batches must be positive.")
 
+    start_time = time.perf_counter()
     paths = _find_scaler_neuron_scaler_paths(model)
     if not paths:
+        logger.info(
+            "Delay estimated: matched_paths=0 batches=0 raw_delay_start=0 delay_start=0 readout_clamped=False elapsed_ms={:.3f}",
+            (time.perf_counter() - start_time) * 1000.0,
+        )
         return 0
 
     original_device = None
@@ -225,6 +216,7 @@ def estimate_delay_start(
     original_training_modes = {module: module.training for module in model.modules()}
     ratios: Dict[BaseNode, List[float]] = {module: [] for module, _ in paths}
     handles = []
+    processed_batches = 0
 
     def make_hook(module: BaseNode, post_scaler: Scaler):
         def hook(_module, inputs, _output):
@@ -245,7 +237,7 @@ def estimate_delay_start(
         for module, post_scaler in paths:
             handles.append(module.register_forward_hook(make_hook(module, post_scaler)))
 
-        _reset_snn(model)
+        reset_net(model)
         with torch.no_grad():
             for batch_idx, batch in enumerate(dataloader):
                 if batch_idx >= num_batches:
@@ -253,8 +245,9 @@ def estimate_delay_start(
                 x = _extract_batch_input(batch)
                 if not isinstance(x, torch.Tensor):
                     raise TypeError("The extracted model input must be a tensor.")
-                _reset_snn(model)
+                reset_net(model)
                 model(x.to(device, non_blocking=True))
+                processed_batches += 1
         delay = 0.0
         for values in ratios.values():
             if values:
@@ -269,14 +262,14 @@ def estimate_delay_start(
     finally:
         for handle in handles:
             handle.remove()
-        _reset_snn(model)
+        reset_net(model)
         if original_device is not None:
             model.to(original_device)
         for module, training in original_training_modes.items():
             module.training = training
 
-    delay_start = math.ceil(delay) if math.isfinite(delay) else time_steps
-    if time_steps < delay_start + _MIN_READOUT_STEPS:
+    raw_delay_start = math.ceil(delay) if math.isfinite(delay) else time_steps
+    if time_steps < raw_delay_start + _MIN_READOUT_STEPS:
         warnings.warn(
             "estimate_delay_start: time_steps is too small to keep "
             f"{_MIN_READOUT_STEPS} readout steps after the estimated delay; "
@@ -284,5 +277,25 @@ def estimate_delay_start(
             RuntimeWarning,
             stacklevel=2,
         )
-        return max(time_steps - _MIN_READOUT_STEPS - 1, 0)
-    return min(delay_start, time_steps - 1)
+        result = max(time_steps - _MIN_READOUT_STEPS - 1, 0)
+        logger.info(
+            "Delay estimated: matched_paths={} batches={} raw_delay_start={} delay_start={} readout_clamped={} elapsed_ms={:.3f}",
+            len(paths),
+            processed_batches,
+            raw_delay_start,
+            result,
+            True,
+            (time.perf_counter() - start_time) * 1000.0,
+        )
+        return result
+    result = min(raw_delay_start, time_steps - 1)
+    logger.info(
+        "Delay estimated: matched_paths={} batches={} raw_delay_start={} delay_start={} readout_clamped={} elapsed_ms={:.3f}",
+        len(paths),
+        processed_batches,
+        raw_delay_start,
+        result,
+        False,
+        (time.perf_counter() - start_time) * 1000.0,
+    )
+    return result

{spikingjelly-2.0.0.dev0 → spikingjelly-2.0.0.dev1}/spikingjelly/activation_based/ann2snn/examples/bert_sst2_transformer_td_equivalent.py

@@ -59,7 +59,8 @@ class FXFriendlyBertSelfAttention(nn.Module):
 
     def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
         x = x.view(
-            *x.size()[:-1],
+            x.shape[0],
+            x.shape[1],
             self.num_attention_heads,
             self.attention_head_size,
         )
@@ -80,7 +81,8 @@ class FXFriendlyBertSelfAttention(nn.Module):
         attention_probs = self.dropout(attention_probs)
         context_layer = torch.matmul(attention_probs, value_layer)
         context_layer = context_layer.permute(0, 2, 1, 3).reshape(
-            *hidden_states.size()[:-1],
+            hidden_states.shape[0],
+            hidden_states.shape[1],
             self.all_head_size,
         )
         return context_layer

{spikingjelly-2.0.0.dev0 → spikingjelly-2.0.0.dev1}/spikingjelly/activation_based/ann2snn/examples/cnn_mnist.py

@@ -9,7 +9,7 @@ import torch
 import torchvision
 from tqdm import tqdm
 
-from spikingjelly.activation_based import ann2snn
+from spikingjelly.activation_based import ann2snn, functional
 from spikingjelly.activation_based.ann2snn.sample_models import mnist_cnn
 
 
@@ -47,7 +47,6 @@ def val(net, device, data_loader, T=None):
     total = 0.0
     if T is not None:
         corrects = np.zeros(T)
-        reset_modules = [m for m in net.modules() if hasattr(m, "reset")]
     with torch.no_grad():
         for batch, (img, label) in enumerate(tqdm(data_loader)):
             img = img.to(device, non_blocking=True)
@@ -56,8 +55,7 @@ def val(net, device, data_loader, T=None):
                 out = net(img)
                 correct += (out.argmax(dim=1) == label).float().sum().item()
             else:
-                for m in reset_modules:
-                    m.reset()
+                functional.reset_net(net)
                 out = None
                 for t in range(T):
                     step = net(img)
@@ -157,8 +155,6 @@ def run_recipe_comparison(
     ltb_accs = convert_and_eval(
         ann2snn.LocalThresholdBalancingRecipe(
             dataloader=calibration_data_loader,
-            time_steps=time_steps,
-            mode="99.9%",
         ),
         device,
         model,
@@ -293,9 +289,6 @@ def main(args):
         transform=torchvision.transforms.ToTensor(),
         download=True,
     )
-    train_data_loader = torch.utils.data.DataLoader(
-        dataset=train_data_dataset, batch_size=batch_size, shuffle=True, drop_last=False
-    )
     calibration_data_loader = torch.utils.data.DataLoader(
         dataset=train_data_dataset,
         batch_size=batch_size,
@@ -315,21 +308,24 @@ def main(args):
         drop_last=False,
     )
 
-    # loss_function = nn.CrossEntropyLoss()
-    # optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=5e-4)
-    # for epoch in range(epochs):
+    # To train the ANN checkpoint locally:
+    # model = mnist_cnn.CNN().to(device)
+    # loss_function = torch.nn.CrossEntropyLoss()
+    # optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)
+    # train_data_loader = torch.utils.data.DataLoader(
+    #     train_data_dataset, batch_size=batch_size, shuffle=True
+    # )
+    # for epoch in range(10):
     #     model.train()
-    #     for (img, label) in train_data_loader:
+    #     for img, label in train_data_loader:
     #         optimizer.zero_grad()
     #         out = model(img.to(device))
     #         loss = loss_function(out, label.to(device))
     #         loss.backward()
     #         optimizer.step()
-    #     torch.save(model.state_dict(), 'SJ-mnist-cnn_model-sample.pth')
-    #     print('Epoch: %d' % epoch)
-    #     acc = val(model, device, train_data_loader)
-    #     print('Validating Accuracy: %.3f' % (acc))
-    #     print()
+    #     torch.save(model.state_dict(), "SJ-mnist-cnn_model-sample.pth")
+    #     print("Epoch: %d" % epoch)
+    #     print("Validating Accuracy: %.3f" % val(model, device, train_data_loader))
 
     if args.plot_mode_sweep:
         run_legacy_mode_sweep(