Skip to content

Commit

Permalink
change to private funcs, add manual optimizer change test
Browse files Browse the repository at this point in the history
  • Loading branch information
@AlbinSou
AlbinSou committed Mar 25, 2024
1 parent 2107dbe commit 7f3a267
Show file tree
Hide file tree
Showing 3 changed files with 105 additions and 28 deletions.
31 changes: 17 additions & 14 deletions avalanche/models/dynamic_optimizers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
colors[None] = colors["END"]


def map_optimized_params(optimizer, parameters, old_params=None):
def _map_optimized_params(optimizer, parameters, old_params=None):
"""
Establishes a mapping between a list of named parameters and the parameters
that are in the optimizer, additionally,
Expand Down Expand Up @@ -105,8 +105,8 @@ def map_optimized_params(optimizer, parameters, old_params=None):
)


def build_tree_from_name_groups(name_groups):
root = TreeNode("") # Root node
def _build_tree_from_name_groups(name_groups):
root = _TreeNode("") # Root node
node_mapping = {}

# Iterate through each string in the list
Expand All @@ -117,7 +117,7 @@ def build_tree_from_name_groups(name_groups):
# Traverse the tree and construct nodes for each component
for component in components:
if component not in current_node.children:
current_node.children[component] = TreeNode(
current_node.children[component] = _TreeNode(
component, parent=current_node
)
current_node = current_node.children[component]
Expand All @@ -135,7 +135,7 @@ def build_tree_from_name_groups(name_groups):
return root, node_mapping


def print_group_information(node, prefix=""):
def _print_group_information(node, prefix=""):
# Print the groups for the current node

if len(node.groups) == 1:
Expand All @@ -150,10 +150,10 @@ def print_group_information(node, prefix=""):

# Recursively print group information for children nodes
for child_name, child_node in node.children.items():
print_group_information(child_node, prefix + " ")
_print_group_information(child_node, prefix + " ")


class ParameterGroupStructure:
class _ParameterGroupStructure:
"""
Structure used for the resolution of unknown parameter groups,
stores parameters as a tree and propagates parameter groups from leaves of
Expand All @@ -162,15 +162,15 @@ class ParameterGroupStructure:

def __init__(self, name_groups, verbose=False):
# Here we rebuild the tree
self.root, self.node_mapping = build_tree_from_name_groups(name_groups)
self.root, self.node_mapping = _build_tree_from_name_groups(name_groups)
if verbose:
print_group_information(self.root)
_print_group_information(self.root)

def __getitem__(self, name):
return self.node_mapping[name]


class TreeNode:
class _TreeNode:
def __init__(self, name, parent=None):
self.name = name
self.children = {}
Expand Down Expand Up @@ -293,7 +293,7 @@ def update_optimizer(
changed_parameters,
new_parameters,
not_found_in_parameters,
) = map_optimized_params(optimizer, new_params, old_params=optimized_params)
) = _map_optimized_params(optimizer, new_params, old_params=optimized_params)

# Change reference to already existing parameters
# i.e growing IncrementalClassifier
Expand All @@ -320,7 +320,7 @@ def update_optimizer(

# Add newly added parameters (i.e Multitask, PNN)

param_structure = ParameterGroupStructure(group_mapping, verbose=verbose)
param_structure = _ParameterGroupStructure(group_mapping, verbose=verbose)

# New parameters
for key in new_parameters:
Expand All @@ -336,13 +336,16 @@ def update_optimizer(
return new_params


@deprecated(
0.6,
"parameters have to be added manually to the optimizer in an existing or a new parameter group",
)
def add_new_params_to_optimizer(optimizer, new_params):
"""Add new parameters to the trainable parameters.
:param new_params: list of trainable parameters
"""
# optimizer.add_param_group({"params": new_params})
optimizer.param_groups[0]["params"].append(new_params)
optimizer.add_param_group({"params": new_params})


__all__ = ["add_new_params_to_optimizer", "reset_optimizer", "update_optimizer"]
13 changes: 8 additions & 5 deletions examples/multihead_param_groups.py → examples/optimizer_param_groups.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,19 @@
# Copyrights licensed under the MIT License. #
# See the accompanying LICENSE file for terms. #
# #
# Date: 01-12-2020 #
# Author(s): Andrea Cossu #
# Date: 25-03-2024 #
# Author(s): Albin Soutif #
# E-mail: [email protected] #
# Website: avalanche.continualai.org #
################################################################################

"""
This example trains a Multi-head model on Split MNIST with Elastich Weight
This example trains a Multi-head model on Split MNIST with Elastic Weight
Consolidation. Each experience has a different task label, which is used at test
time to select the appropriate head.
time to select the appropriate head. Additionally, it assigns different parameter groups
to the classifier and the backbone, assigning lower learning rate to
the backbone than to the classifier. When the multihead classifier grows,
new parameters are automatically assigned to the corresponding parameter group
"""

import argparse
Expand Down Expand Up @@ -78,7 +81,7 @@ def main(args):

# train and test loop
for train_task in train_stream:
strategy.train(train_task, num_workers=4)
strategy.train(train_task, num_workers=4, verbose=True)
strategy.eval(test_stream, num_workers=4)


Expand Down
89 changes: 80 additions & 9 deletions tests/test_optimizer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
from avalanche.models.cosine_layer import CosineLinear, SplitCosineLinear
from avalanche.models.dynamic_optimizers import (
add_new_params_to_optimizer,
map_optimized_params,
update_optimizer,
)
from avalanche.models.pytorchcv_wrapper import densenet, get_model, pyramidnet, resnet
Expand Down Expand Up @@ -85,12 +84,6 @@ def _is_param_in_optimizer_group(self, param, optimizer):
return None

def load_benchmark(self, use_task_labels=False):
"""
Returns a NC benchmark from a fake dataset of 10 classes, 5 experiences,
2 classes per experience.
:param fast_test: if True loads fake data, MNIST otherwise.
"""
return get_fast_benchmark(use_task_labels=use_task_labels)

def init_scenario(self, multi_task=False):
Expand All @@ -104,9 +97,9 @@ def test_optimizer_update(self):
optimizer = SGD(model.parameters(), lr=1e-3)
strategy = Naive(model=model, optimizer=optimizer)

# check add_param_group
# check add new parameter
p = torch.nn.Parameter(torch.zeros(10, 10))
add_new_params_to_optimizer(optimizer, p)
optimizer.param_groups[0]["params"].append(p)
assert self._is_param_in_optimizer(p, strategy.optimizer)

# check new_param is in optimizer
Expand All @@ -125,6 +118,10 @@ def test_optimizer_update(self):
self.assertFalse(self._is_param_in_optimizer(p, strategy.optimizer))

def test_optimizers(self):
"""
Run a series of tests on various pytorch optimizers
"""

# SIT scenario
model, criterion, benchmark = self.init_scenario(multi_task=True)
for optimizer in self._iterate_optimizers(
Expand All @@ -143,6 +140,10 @@ def test_optimizers(self):
self._test_optimizer_state(strategy)

def test_optimizer_groups_clf_til(self):
"""
Tests the automatic assignation of new
MultiHead parameters to the optimizer
"""
model, criterion, benchmark = self.init_scenario(multi_task=True)

g1 = []
Expand Down Expand Up @@ -180,6 +181,10 @@ def test_optimizer_groups_clf_til(self):
)

def test_optimizer_groups_clf_cil(self):
"""
Tests the automatic assignation of new
IncrementalClassifier parameters to the optimizer
"""
model, criterion, benchmark = self.init_scenario(multi_task=False)

g1 = []
Expand Down Expand Up @@ -216,7 +221,73 @@ def test_optimizer_groups_clf_cil(self):
self._is_param_in_optimizer_group(p, strategy.optimizer), 1
)

def test_optimizer_groups_manual_addition(self):
"""
Tests the manual addition of a new parameter group
mixed with existing MultiHeadClassifier
"""
model, criterion, benchmark = self.init_scenario(multi_task=True)

g1 = []
g2 = []
for n, p in model.named_parameters():
if "classifier" in n:
g1.append(p)
else:
g2.append(p)

optimizer = SGD([{"params": g1, "lr": 0.1}, {"params": g2, "lr": 0.05}])

strategy = Naive(
model=model,
optimizer=optimizer,
criterion=criterion,
train_mb_size=64,
device=self.device,
eval_mb_size=50,
train_epochs=2,
)

experience_0 = benchmark.train_stream[0]
experience_1 = benchmark.train_stream[1]

strategy.train(experience_0)

# Add some new parameter and assign it manually to param group
model.new_module1 = torch.nn.Linear(10, 10)
model.new_module2 = torch.nn.Linear(10, 10)
strategy.optimizer.param_groups[1]["params"] += list(
model.new_module1.parameters()
)
strategy.optimizer.add_param_group(
{"params": list(model.new_module2.parameters()), "lr": 0.001}
)

# Also add one but to a new param group

strategy.train(experience_1)

for n, p in model.named_parameters():
assert self._is_param_in_optimizer(p, strategy.optimizer)
if "classifier" in n:
self.assertEqual(
self._is_param_in_optimizer_group(p, strategy.optimizer), 0
)
elif "new_module2" in n:
self.assertEqual(
self._is_param_in_optimizer_group(p, strategy.optimizer), 2
)
else:
self.assertEqual(
self._is_param_in_optimizer_group(p, strategy.optimizer), 1
)

def test_optimizer_groups_rename(self):
"""
Tests the correct reassignation to
existing parameter groups after
parameter renaming
"""
model, criterion, benchmark = self.init_scenario(multi_task=False)

g1 = []
Expand Down

0 comments on commit 7f3a267

Please sign in to comment.