From 0785eae552de0dcfc13d54af73cac78a5a56b4b1 Mon Sep 17 00:00:00 2001 From: yoshikisd Date: Fri, 1 Aug 2025 21:37:05 +0000 Subject: [PATCH] Created Reconstructors class to replace optimization methods in CDIModel --- src/cdtools/__init__.py | 4 +- src/cdtools/reconstructors/__init__.py | 16 ++ src/cdtools/reconstructors/adam.py | 160 ++++++++++++ src/cdtools/reconstructors/base.py | 325 +++++++++++++++++++++++++ src/cdtools/reconstructors/lbfgs.py | 134 ++++++++++ src/cdtools/reconstructors/sgd.py | 157 ++++++++++++ 6 files changed, 794 insertions(+), 2 deletions(-) create mode 100644 src/cdtools/reconstructors/__init__.py create mode 100644 src/cdtools/reconstructors/adam.py create mode 100644 src/cdtools/reconstructors/base.py create mode 100644 src/cdtools/reconstructors/lbfgs.py create mode 100644 src/cdtools/reconstructors/sgd.py diff --git a/src/cdtools/__init__.py b/src/cdtools/__init__.py index 9684207..9132209 100644 --- a/src/cdtools/__init__.py +++ b/src/cdtools/__init__.py @@ -4,9 +4,9 @@ import warnings warnings.filterwarnings("ignore", message='To copy construct from a tensor, ') -__all__ = ['tools', 'datasets', 'models'] +__all__ = ['tools', 'datasets', 'models', 'reconstructors'] from cdtools import tools from cdtools import datasets from cdtools import models - +from cdtools import reconstructors diff --git a/src/cdtools/reconstructors/__init__.py b/src/cdtools/reconstructors/__init__.py new file mode 100644 index 0000000..84b96ab --- /dev/null +++ b/src/cdtools/reconstructors/__init__.py @@ -0,0 +1,16 @@ +"""This module contains optimizers for performing reconstructions + +""" + +# We define __all__ to be sure that import * only imports what we want +__all__ = [ + 'Reconstructor', + 'Adam', + 'LBFGS', + 'SGD' +] + +from cdtools.reconstructors.base import Reconstructor +from cdtools.reconstructors.adam import Adam +from cdtools.reconstructors.lbfgs import LBFGS +from cdtools.reconstructors.sgd import SGD diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py new file mode 100644 index 0000000..5a489a0 --- /dev/null +++ b/src/cdtools/reconstructors/adam.py @@ -0,0 +1,160 @@ +"""This module contains the Adam Reconstructor subclass for performing +optimization ('reconstructions') on ptychographic/CDI models using +the Adam optimizer. + +The Reconstructor class is designed to resemble so-called +'Trainer' classes that (in the language of the AI/ML folks) handles +the 'training' of a model given some dataset and optimizer. +""" +import torch as t +from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset +from cdtools.models import CDIModel +from typing import Tuple, List, Union +from cdtools.reconstructors import Reconstructor + +__all__ = ['Adam'] + + +class Adam(Reconstructor): + """ + The Adam Reconstructor subclass handles the optimization ('reconstruction') + of ptychographic models and datasets using the Adam optimizer. + + Parameters + ---------- + model: CDIModel + Model for CDI/ptychography reconstruction. + dataset: Ptycho2DDataset + The dataset to reconstruct against. + subset : list(int) or int + Optional, a pattern index or list of pattern indices to use. + schedule : bool + Optional, create a learning rate scheduler + (torch.optim.lr_scheduler._LRScheduler). + + Important attributes: + - **model** -- Always points to the core model used. + - **optimizer** -- This class by default uses `torch.optim.Adam` to perform + optimizations. + - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the + `optimize` method. + - **data_loader** -- A torch.utils.data.DataLoader that is defined by + calling the `setup_dataloader` method. + """ + def __init__(self, + model: CDIModel, + dataset: Ptycho2DDataset, + subset: List[int] = None): + + super().__init__(model, dataset, subset) + + # Define the optimizer for use in this subclass + self.optimizer = t.optim.Adam(self.model.parameters()) + + def adjust_optimizer(self, + lr: int = 0.005, + betas: Tuple[float] = (0.9, 0.999), + amsgrad: bool = False): + """ + Change hyperparameters for the utilized optimizer. + + Parameters + ---------- + lr : float + Optional, The learning rate (alpha) to use. Default is 0.005. 0.05 + is typically the highest possible value with any chance of being + stable. + betas : tuple + Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999). + amsgrad : bool + Optional, whether to use the AMSGrad variant of this algorithm. + """ + for param_group in self.optimizer.param_groups: + param_group['lr'] = lr + param_group['betas'] = betas + param_group['amsgrad'] = amsgrad + + def optimize(self, + iterations: int, + batch_size: int = 15, + lr: float = 0.005, + betas: Tuple[float] = (0.9, 0.999), + schedule: bool = False, + amsgrad: bool = False, + regularization_factor: Union[float, List[float]] = None, + thread: bool = True, + calculation_width: int = 10, + shuffle: bool = True): + """ + Runs a round of reconstruction using the Adam optimizer + + Formerly `CDIModel.Adam_optimize` + + This calls the Reconstructor.optimize superclass method + (formerly `CDIModel.AD_optimize`) to run a round of reconstruction + once the dataloader and optimizer hyperparameters have been + set up. + + Parameters + ---------- + iterations : int + How many epochs of the algorithm to run. + batch_size : int + Optional, the size of the minibatches to use. + lr : float + Optional, The learning rate (alpha) to use. Default is 0.005. 0.05 + is typically the highest possible value with any chance of being + stable. + betas : tuple + Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999). + schedule : bool + Optional, create a learning rate scheduler + (torch.optim.lr_scheduler._LRScheduler). + amsgrad : bool + Optional, whether to use the AMSGrad variant of this algorithm. + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of + parameters to pass the regularizer method. + thread : bool + Default True, whether to run the computation in a separate thread + to allow interaction with plots during computation. + calculation_width : int + Default 10, how many translations to pass through at once for each + round of gradient accumulation. Does not affect the result, only + the calculation speed. + shuffle : bool + Optional, enable/disable shuffling of the dataset. This option + is intended for diagnostic purposes and should be left as True. + """ + # Update the training history + self.model.training_history += ( + f'Planning {iterations} epochs of Adam, with a learning rate = ' + f'{lr}, batch size = {batch_size}, regularization_factor = ' + f'{regularization_factor}, and schedule = {schedule}.\n' + ) + + # 1) The subset statement is contained in Reconstructor.__init__ + + # 2) Set up / re-initialize the data laoder + self.setup_dataloader(batch_size=batch_size, shuffle=shuffle) + + # 3) The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer + self.adjust_optimizer(lr=lr, + betas=betas, + amsgrad=amsgrad) + + # 4) Set up the scheduler + if schedule: + self.scheduler = \ + t.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, + factor=0.2, + threshold=1e-9) + else: + self.scheduler = None + + # 5) This is analagous to making a call to CDIModel.AD_optimize + return super(Adam, self).optimize(iterations, + regularization_factor, + thread, + calculation_width) diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py new file mode 100644 index 0000000..65b2771 --- /dev/null +++ b/src/cdtools/reconstructors/base.py @@ -0,0 +1,325 @@ +"""This module contains the base Reconstructor class for performing +optimization ('reconstructions') on ptychographic/CDI models. + +The Reconstructor class is designed to resemble so-called +'Trainer' classes that (in the language of the AI/ML folks) handles +the 'training' of a model given some dataset and optimizer. + +The subclasses of Reconstructor are required to implement +their own data loaders and optimizer adjusters +""" + +import torch as t +from torch.utils import data as td +import threading +import queue +import time +from cdtools.datasets import CDataset +from cdtools.models import CDIModel +from typing import List, Union + +__all__ = ['Reconstructor'] + + +class Reconstructor: + """ + Reconstructor handles the optimization ('reconstruction') of ptychographic + models given a CDIModel (or subclass) and corresponding CDataset. + + This is a base model that defines all functions Reconstructor subclasses + must implement. + + Parameters + ---------- + model: CDIModel + Model for CDI/ptychography reconstruction + dataset: CDataset + The dataset to reconstruct against + subset : list(int) or int + Optional, a pattern index or list of pattern indices to use + + Important attributes: + - **model** -- Always points to the core model used. + - **optimizer** -- A `torch.optim.Optimizer` that must be defined when + initializing the Reconstructor subclass. + - **scheduler** -- A `torch.optim.lr_scheduler` that may be defined during + the `optimize` method. + - **data_loader** -- A torch.utils.data.DataLoader that is defined by + calling the `setup_dataloader` method. + """ + def __init__(self, + model: CDIModel, + dataset: CDataset, + subset: Union[int, List[int]] = None): + # Store parameters as attributes of Reconstructor + self.subset = subset + + # Initialize attributes that must be defined by the subclasses + self.optimizer = None + self.scheduler = None + self.data_loader = None + + # Store the original model + self.model = model + + # Store the dataset + if subset is not None: + # if subset is just one pattern, turn into a list for convenience + if isinstance(subset, int): + subset = [subset] + dataset = td.Subset(dataset, subset) + self.dataset = dataset + + def setup_dataloader(self, + batch_size: int = None, + shuffle: bool = True): + """ + Sets up / re-initializes the dataloader. + + Parameters + ---------- + batch_size : int + Optional, the size of the minibatches to use + shuffle : bool + Optional, enable/disable shuffling of the dataset. This option + is intended for diagnostic purposes and should be left as True. + """ + if batch_size is not None: + self.data_loader = td.DataLoader(self.dataset, + batch_size=batch_size, + shuffle=shuffle) + else: + self.data_loader = td.Dataloader(self.dataset) + + def adjust_optimizer(self, **kwargs): + """ + Change hyperparameters for the utilized optimizer. + + For each optimizer, the keyword arguments should be manually defined + as parameters. + """ + raise NotImplementedError() + + def _run_epoch(self, + stop_event: threading.Event = None, + regularization_factor: Union[float, List[float]] = None, + calculation_width: int = 10): + """ + Runs one full epoch of the reconstruction. Intended to be called + by Reconstructor.optimize. + + Parameters + ---------- + stop_event : threading.Event + Default None, causes the reconstruction to stop when an exception + occurs in Optimizer.optimize. + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of + parameters to pass the regularizer method + calculation_width : int + Default 10, how many translations to pass through at once for each + round of gradient accumulation. This does not affect the result, + but may affect the calculation speed. + + Returns + ------ + loss : float + The summed loss over the latest epoch, divided by the total + diffraction pattern intensity + """ + + # Initialize some tracking variables + normalization = 0 + loss = 0 + N = 0 + t0 = time.time() + + # The data loader is responsible for setting the minibatch + # size, so each set is a minibatch + for inputs, patterns in self.data_loader: + normalization += t.sum(patterns).cpu().numpy() + N += 1 + + def closure(): + self.optimizer.zero_grad() + + # We further break up the minibatch into a set of chunks. + # This lets us use larger minibatches than can fit + # on the GPU at once, while still doing batch processing + # for efficiency + input_chunks = [[inp[i:i + calculation_width] + for inp in inputs] + for i in range(0, len(inputs[0]), + calculation_width)] + pattern_chunks = [patterns[i:i + calculation_width] + for i in range(0, len(inputs[0]), + calculation_width)] + + total_loss = 0 + + for inp, pats in zip(input_chunks, pattern_chunks): + # This check allows for graceful exit when threading + if stop_event is not None and stop_event.is_set(): + exit() + + # Run the simulation + sim_patterns = self.model.forward(*inp) + + # Calculate the loss + if hasattr(self, 'mask'): + loss = self.model.loss(pats, + sim_patterns, + mask=self.model.mask) + else: + loss = self.model.loss(pats, + sim_patterns) + + # And accumulate the gradients + loss.backward() + + # Normalize the accumulating total loss + total_loss += loss.detach() // self.model.world_size + + # If we have a regularizer, we can calculate it separately, + # and the gradients will add to the minibatch gradient + if regularization_factor is not None \ + and hasattr(self.model, 'regularizer'): + + loss = self.model.regularizer(regularization_factor) + loss.backward() + + return total_loss + + # This takes the step for this minibatch + loss += self.optimizer.step(closure).detach().cpu().numpy() + + loss /= normalization + + # We step the scheduler after the full epoch + if self.scheduler is not None: + self.scheduler.step(loss) + + self.model.loss_history.append(loss) + self.model.epoch = len(self.model.loss_history) + self.model.latest_iteration_time = time.time() - t0 + self.model.training_history += self.model.report() + '\n' + return loss + + def optimize(self, + iterations: int, + regularization_factor: Union[float, List[float]] = None, + thread: bool = True, + calculation_width: int = 10): + """ + Runs a round of reconstruction using the provided optimizer + + Formerly CDIModel.AD_optimize + + This is the basic automatic differentiation reconstruction tool + which all the other, algorithm-specific tools, use. It is a + generator which yields the average loss each epoch, ending after + the specified number of iterations. + + By default, the computation will be run in a separate thread. This + is done to enable live plotting with matplotlib during a + reconstruction. + + If the computation was done in the main thread, this would freeze + the plots. This behavior can be turned off by setting the keyword + argument 'thread' to False. + + Parameters + ---------- + iterations : int + How many epochs of the algorithm to run. + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of + parameters to pass the regularizer method. + thread : bool + Default True, whether to run the computation in a separate thread + to allow interaction with plots during computation. + calculation_width : int + Default 10, how many translations to pass through at once for each + round of gradient accumulation. This does not affect the result, + but may affect the calculation speed. + + Yields + ------ + loss : float + The summed loss over the latest epoch, divided by the total + diffraction pattern intensity. + """ + + # We store the current optimizer as a model parameter so that + # it can be saved and loaded for checkpointing + self.current_optimizer = self.optimizer + + # If we don't want to run in a different thread, this is easy + if not thread: + for it in range(iterations): + if self.model.skip_computation(): + self.epoch = self.epoch + 1 + if len(self.model.loss_history) >= 1: + yield self.model.loss_history[-1] + else: + yield float('nan') + continue + + yield self._run_epoch(regularization_factor=regularization_factor, # noqa + calculation_width=calculation_width) + + # But if we do want to thread, it's annoying: + else: + # Here we set up the communication with the computation thread + result_queue = queue.Queue() + stop_event = threading.Event() + + def target(): + try: + result_queue.put( + self._run_epoch(stop_event=stop_event, + regularization_factor=regularization_factor, # noqa + calculation_width=calculation_width)) + except Exception as e: + # If something bad happens, put the exception into the + # result queue + result_queue.put(e) + + # And this actually starts and monitors the thread + for it in range(iterations): + if self.model.skip_computation(): + self.model.epoch = self.model.epoch + 1 + if len(self.model.loss_history) >= 1: + yield self.model.loss_history[-1] + else: + yield float('nan') + continue + + calc = threading.Thread(target=target, + name='calculator', + daemon=True) + try: + calc.start() + while calc.is_alive(): + if hasattr(self.model, 'figs'): + self.model.figs[0].canvas.start_event_loop(0.01) + else: + calc.join() + + except KeyboardInterrupt as e: + stop_event.set() + print('\nAsking execution thread to stop cleanly - ' + + 'please be patient.') + calc.join() + raise e + + res = result_queue.get() + + # If something went wrong in the thead, we'll get an exception + if isinstance(res, Exception): + raise res + + yield res + + # And finally, we unset the current optimizer: + self.current_optimizer = None diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py new file mode 100644 index 0000000..0b51dfd --- /dev/null +++ b/src/cdtools/reconstructors/lbfgs.py @@ -0,0 +1,134 @@ +"""This module contains the LBFGS Reconstructor subclass for performing +optimization ('reconstructions') on ptychographic/CDI models using +the LBFGS optimizer. + +The Reconstructor class is designed to resemble so-called +'Trainer' classes that (in the language of the AI/ML folks) handles +the 'training' of a model given some dataset and optimizer. +""" +import torch as t +from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset +from cdtools.models import CDIModel +from typing import List, Union +from cdtools.reconstructors import Reconstructor + +__all__ = ['LBFGS'] + + +class LBFGS(Reconstructor): + """ + The LBFGS Reconstructor subclass handles the optimization + ('reconstruction') of ptychographic models and datasets using the LBFGS + optimizer. + + Parameters + ---------- + model: CDIModel + Model for CDI/ptychography reconstruction. + dataset: Ptycho2DDataset + The dataset to reconstruct against. + subset : list(int) or int + Optional, a pattern index or list of pattern indices to use. + schedule : bool + Optional, create a learning rate scheduler + (torch.optim.lr_scheduler._LRScheduler). + + Important attributes: + - **model** -- Always points to the core model used. + - **optimizer** -- This class by default uses `torch.optim.LBFGS` to + perform optimizations. + - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during + the `optimize` method. + - **data_loader** -- A torch.utils.data.DataLoader that is defined by + calling the `setup_dataloader` method. + """ + def __init__(self, + model: CDIModel, + dataset: Ptycho2DDataset, + subset: List[int] = None): + + super().__init__(model, dataset, subset) + + # Define the optimizer for use in this subclass + self.optimizer = t.optim.LBFGS(self.model.parameters()) + + def adjust_optimizer(self, + lr: int = 0.005, + history_size: int = 2, + line_search_fn: str = None): + """ + Change hyperparameters for the utilized optimizer. + + Parameters + ---------- + lr : float + Optional, The learning rate (alpha) to use. Default is 0.005. 0.05 + is typically the highest possible value with any chance of being + stable. + history_size : int + Optional, the length of the history to use. + line_search_fn : str + Optional, either `strong_wolfe` or None + """ + for param_group in self.optimizer.param_groups: + param_group['lr'] = lr + param_group['history_size'] = history_size + param_group['line_search_fn'] = line_search_fn + + def optimize(self, + iterations: int, + lr: float = 0.1, + history_size: int = 2, + regularization_factor: Union[float, List[float]] = None, + thread: bool = True, + calculation_width: int = 10, + line_search_fn: str = None): + """ + Runs a round of reconstruction using the LBFGS optimizer + + Formerly `CDIModel.LBFGS_optimize` + + This algorithm is often less stable that Adam, however in certain + situations or geometries it can be shockingly efficient. Like all + the other optimization routines, it is defined as a generator + function which yields the average loss each epoch. + + NOTE: There is no batch size, because it is a usually a bad idea to use + LBFGS on anything but all the data at onece + + Parameters + ---------- + iterations : int + How many epochs of the algorithm to run. + lr : float + Optional, The learning rate (alpha) to use. Default is 0.1. + history_size : int + Optional, the length of the history to use. + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of + parameters to pass the regularizer method. + thread : bool + Default True, whether to run the computation in a separate thread + to allow interaction with plots during computation. + calculation_width : int + Default 10, how many translations to pass through at once for each + round of gradient accumulation. Does not affect the result, only + the calculation speed. + """ + # 1) The subset statement is contained in Reconstructor.__init__ + + # 2) Set up / re-initialize the data loader. For LBFGS, we load + # all the data at once. + self.setup_dataloader(batch_size=len(self.dataset)) + + # 3) The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer + self.adjust_optimizer(lr=lr, + history_size=history_size, + line_search_fn=line_search_fn) + + # 4) This is analagous to making a call to CDIModel.AD_optimize + return super(LBFGS, self).optimize(iterations, + regularization_factor, + thread, + calculation_width) diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py new file mode 100644 index 0000000..f2dd7b0 --- /dev/null +++ b/src/cdtools/reconstructors/sgd.py @@ -0,0 +1,157 @@ +"""This module contains the SGD Reconstructor subclass for performing +optimization ('reconstructions') on ptychographic/CDI models using +stochastic gradient descent. + +The Reconstructor class is designed to resemble so-called +'Trainer' classes that (in the language of the AI/ML folks) handles +the 'training' of a model given some dataset and optimizer. +""" +import torch as t +from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset +from cdtools.models import CDIModel +from typing import List, Union +from cdtools.reconstructors import Reconstructor + +__all__ = ['SGD'] + + +class SGD(Reconstructor): + """ + The Adam Reconstructor subclass handles the optimization ('reconstruction') + of ptychographic models and datasets using the Adam optimizer. + + Parameters + ---------- + model: CDIModel + Model for CDI/ptychography reconstruction. + dataset: Ptycho2DDataset + The dataset to reconstruct against. + subset : list(int) or int + Optional, a pattern index or list of pattern indices to use. + + Important attributes: + - **model** -- Always points to the core model used. + - **optimizer** -- This class by default uses `torch.optim.Adam` to perform + optimizations. + - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the + `optimize` method. + - **data_loader** -- A torch.utils.data.DataLoader that is defined by + calling the `setup_dataloader` method. + """ + def __init__(self, + model: CDIModel, + dataset: Ptycho2DDataset, + subset: List[int] = None): + + super().__init__(model, dataset, subset) + + # Define the optimizer for use in this subclass + self.optimizer = t.optim.SGD(self.model.parameters()) + + def adjust_optimizer(self, + lr: int = 0.005, + momentum: float = 0, + dampening: float = 0, + weight_decay: float = 0, + nesterov: bool = False): + """ + Change hyperparameters for the utilized optimizer. + + Parameters + ---------- + lr : float + Optional, The learning rate (alpha) to use. Default is 0.005. 0.05 + is typically the highest possible value with any chance of being + stable. + momentum : float + Optional, the length of the history to use. + dampening : float + Optional, dampening for the momentum. + weight_decay : float + Optional, weight decay (L2 penalty). + nesterov : bool + Optional, enables Nesterov momentum. Only applicable when momentum + is non-zero. + """ + for param_group in self.optimizer.param_groups: + param_group['lr'] = lr + param_group['momentum'] = momentum + param_group['dampening'] = dampening + param_group['weight_decay'] = weight_decay + param_group['nesterov'] = nesterov + + def optimize(self, + iterations: int, + batch_size: int = None, + lr: float = 2e-7, + momentum: float = 0, + dampening: float = 0, + weight_decay: float = 0, + nesterov: bool = False, + regularization_factor: Union[float, List[float]] = None, + thread: bool = True, + calculation_width: int = 10, + shuffle: bool = True): + """ + Runs a round of reconstruction using the Adam optimizer + + Formerly `CDIModel.Adam_optimize` + + This calls the Reconstructor.optimize superclass method + (formerly `CDIModel.AD_optimize`) to run a round of reconstruction + once the dataloader and optimizer hyperparameters have been + set up. + + Parameters + ---------- + iterations : int + How many epochs of the algorithm to run. + batch_size : int + Optional, the size of the minibatches to use. + lr : float + Optional, The learning rate to use. The default is 2e-7. + momentum : float + Optional, the length of the history to use. + dampening : float + Optional, dampening for the momentum. + weight_decay : float + Optional, weight decay (L2 penalty). + nesterov : bool + Optional, enables Nesterov momentum. Only applicable when momentum + is non-zero. + regularization_factor : float or list(float) + Optional, if the model has a regularizer defined, the set of + parameters to pass the regularizer method. + thread : bool + Default True, whether to run the computation in a separate thread + to allow interaction with plots during computation. + calculation_width : int + Default 10, how many translations to pass through at once for each + round of gradient accumulation. Does not affect the result, only + the calculation speed. + shuffle : bool + Optional, enable/disable shuffling of the dataset. This option + is intended for diagnostic purposes and should be left as True. + """ + # 1) The subset statement is contained in Reconstructor.__init__ + + # 2) Set up / re-initialize the data laoder + if batch_size is not None: + self.setup_dataloader(batch_size=batch_size, shuffle=shuffle) + else: + # Use default torch dataloader parameters + self.setup_dataloader(batch_size=1, shuffle=False) + + # 3) The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer + self.adjust_optimizer(lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov) + + # 4) This is analagous to making a call to CDIModel.AD_optimize + return super(SGD, self).optimize(iterations, + regularization_factor, + thread, + calculation_width)