<commit_before># ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2019, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import json
import os
from collections import defaultdict
from collections.abc import Iterable
from copy import deepcopy

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as schedulers

from nupic.research.frameworks.dynamic_sparse.networks import NumScheduler
from nupic.torch.modules import update_boost_strength

from .loggers import BaseLogger, SparseLogger
from .modules import SparseModule

__all__ = ["BaseModel", "SparseModel"]


class BaseModel:
    """Base model, with training loops and logging functions."""

    def __init__(self, network=None, config=None):

        defaults = dict(
            optim_alg="SGD",
            learning_rate=0.1,
            momentum=0.9,
            nesterov_momentum=False,
            device="cpu",
            lr_scheduler=False,
            lr_step_size=1,
            pruning_interval=1,
            weight_prune_perc=0,
            grad_prune_perc=0,
            test_noise=False,
            weight_decay=1e-4,
            use_multiple_gpus=False,
            train_batches_per_epoch=np.inf,  # default - don't limit the batches
        )
        defaults.update(config or {})
        self.__dict__.update(defaults)

        # save config to restore the model later
        self.config = config
        self.device = torch.device(self.device)
        if self.use_multiple_gpus:
            network = nn.DataParallel(network)
        self.network = network.to(self.device)
        self.config = deepcopy(config)

    def setup(self, config=None):

        # allow setup to receive a new config
        if config is not None:
            self.__dict__.update(config)

        # init optimizer
        if self.optim_alg == "Adam":
            self.optimizer = optim.Adam(
                self.network.parameters(), lr=self.learning_rate
            )
        elif self.optim_alg == "SGD":
            # added weight decay
            self.optimizer = optim.SGD(
                self.network.parameters(),
                lr=self.learning_rate,
                momentum=self.momentum,
                weight_decay=self.weight_decay,
                nesterov=self.nesterov_momentum,
            )<commit_msg>bug fix to Adam optimizer initialization, added weight decay<commit_after># ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2019, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import json
import os
from collections import defaultdict
from collections.abc import Iterable
from copy import deepcopy

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.lr_scheduler as schedulers

from nupic.research.frameworks.dynamic_sparse.networks import NumScheduler
from nupic.torch.modules import update_boost_strength

from .loggers import BaseLogger, SparseLogger
from .modules import SparseModule

__all__ = ["BaseModel", "SparseModel"]


class BaseModel:
    """Base model, with training loops and logging functions."""

    def __init__(self, network=None, config=None):

        defaults = dict(
            optim_alg="SGD",
            learning_rate=0.1,
            momentum=0.9,
            nesterov_momentum=False,
            device="cpu",
            lr_scheduler=False,
            lr_step_size=1,
            pruning_interval=1,
            weight_prune_perc=0,
            grad_prune_perc=0,
            test_noise=False,
            weight_decay=1e-4,
            use_multiple_gpus=False,
            train_batches_per_epoch=np.inf,  # default - don't limit the batches
        )
        defaults.update(config or {})
        self.__dict__.update(defaults)

        # save config to restore the model later
        self.config = config
        self.device = torch.device(self.device)
        if self.use_multiple_gpus:
            network = nn.DataParallel(network)
        self.network = network.to(self.device)
        self.config = deepcopy(config)

    def setup(self, config=None):

        # allow setup to receive a new config
        if config is not None:
            self.__dict__.update(config)

        # init optimizer
        if self.optim_alg == "Adam":
            self.optimizer = optim.Adam(
                self.network.parameters(),
                lr=self.learning_rate,
                weight_decay=self.weight_decay,
            )
        elif self.optim_alg == "SGD":
            # added weight decay
            self.optimizer = optim.SGD(
                self.network.parameters(),
                lr=self.learning_rate,
                momentum=self.momentum,
                weight_decay=self.weight_decay,
                nesterov=self.nesterov_momentum,
            )
{'fraction_non_alphanumeric': 0.11394921994493729, 'fraction_numerical': 0.004588559192413582, 'mean_word_length': 2.1513253012048192, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3518246', 'n_tokens_mistral': 1754, 'n_tokens_neox': 1691, 'n_words': 553}