super-gradients/tests/unit_tests/zero_weight_decay_on_bias_bn_test.py

import unittest
import torch.nn as nn

from super_gradients.training.utils.optimizer_utils import separate_zero_wd_params_groups_for_optimizer
from super_gradients.training.utils import HpmStruct
from super_gradients.training.models.sg_module import SgModule


class ToySgModule(SgModule):
    """
    Toy Module to test zero of weight decay, support multiple group of parameters.
    """
    CONV_CLASSES = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d}
    CONV_TRANSPOSE_CLASSES = {1: nn.ConvTranspose1d, 2: nn.ConvTranspose2d, 3: nn.ConvTranspose3d}
    BN_CLASSES = {1: nn.BatchNorm1d, 2: nn.BatchNorm2d, 3: nn.BatchNorm3d}

    def __init__(self, input_dimension=2, multiple_param_groups=False, module_groups=False):
        """
        :param input_dimension: input dimension, 1 for 1D, 2 for 2D ...
        :param multiple_param_groups: if True create multiple param groups with different optimizer args.
        """
        super(ToySgModule, self).__init__()
        num_classes = 10

        self.multiple_param_groups = multiple_param_groups
        self.module_groups = module_groups

        self.conv_cls = self.CONV_CLASSES[input_dimension]
        self.bn_cls = self.BN_CLASSES[input_dimension]
        self.conv_transpose_cls = self.CONV_TRANSPOSE_CLASSES[input_dimension]

        self.num_conv = 0
        self.num_bn = 0
        self.num_biases = 0
        self.num_linear = 0

        self.base = nn.Sequential(
            self.conv1(3, 128, 2),
            self.conv1(128, 128, 2, bias=True),
            self.conv_transpose(128, 128),
            self.conv_transpose(128, 128, bias=True),
        )

        self.base_params = (self.num_no_decay_params(), self.num_decay_params())

        self.more_convs = nn.Sequential(
            self.conv1(128, 128, 1),
            self.conv1(128, 128, 2, bias=True),
            self.conv_transpose(128, 128),
        )
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            self.linear(128, 2 * num_classes, bias=False),
            self.linear(2 * num_classes, num_classes, bias=True)
        )

        self.head = nn.Sequential(
            self.more_convs,
            self.avg_pool,
            self.classifier
        )

        self.head_params = (self.num_no_decay_params() - self.base_params[0], self.num_decay_params() - self.base_params[1])

    def conv1(self, ch_in, ch_out, stride, bias=False):
        self.num_conv += 1
        if bias:
            conv = self.conv_cls(ch_in, ch_out, 1, stride=stride, bias=bias)
            self.num_biases += 1
        else:
            conv = nn.Sequential(
                self.conv_cls(ch_in, ch_out, 1, stride=stride, bias=bias),
                self.bn_cls(ch_out),
                nn.ReLU()
            )
            self.num_bn += 1
        return conv

    def conv_transpose(self, ch_in, ch_out, bias=False):
        self.num_conv += 1
        if bias:
            conv = self.conv_transpose_cls(ch_in, ch_out, 2, stride=2, bias=bias)
            self.num_biases += 1
        else:
            conv = nn.Sequential(
                self.conv_transpose_cls(ch_in, ch_out, 2, stride=2, bias=bias),
                self.bn_cls(ch_out),
                nn.ReLU()
            )
            self.num_bn += 1
        return conv

    def linear(self, ch_in, ch_out, bias=False):
        self.num_linear += 1
        if bias:
            self.num_biases += 1
        return nn.Linear(ch_in, ch_out, bias)

    def num_decay_params(self):
        return self.num_conv + self.num_linear

    def num_no_decay_params(self):
        return self.num_biases + 2 * self.num_bn

    def initialize_param_groups(self, lr: float, training_params: HpmStruct) -> list:
        # Example to different learning rates, similar to ShelfNet, in order to create multiple groups.
        if self.multiple_param_groups:
            params_list = [{'named_params': self.base.named_parameters(), 'lr': lr},
                           {'named_params': self.head.named_parameters(), 'lr': lr * 10}]

            return params_list

        return super().initialize_param_groups(lr, training_params)


class ZeroWdForBnBiasTest(unittest.TestCase):
    """
    Testing if the optimizer parameters are divided into two groups with one being with weight_decay = 0
    """

    def setUp(self):
        # Define Parameters
        self.weight_decay = 0.01
        self.lr = 0.1
        # input dimension beside batch and channels, i.e 2 for vision, 1 for audio, 3 for point cloud.
        self.input_dimensions = (1, 2, 3)
        self.train_params_zero_wd = {"initial_lr": self.lr,
                                     "optimizer": "SGD",
                                     "optimizer_params": {"weight_decay": self.weight_decay, "momentum": 0.9}}

    def test_zero_wd_one_group(self):
        """
        test that one group of parameters are separated to weight_decay_params and without.
        """
        for input_dim in self.input_dimensions:
            net = ToySgModule(input_dimension=input_dim)
            train_params = HpmStruct(**self.train_params_zero_wd)

            optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer(
                net,
                net.initialize_param_groups(self.lr, train_params),
                self.weight_decay
            )
            self.assertEqual(len(optimizer_params_groups), 2, msg="Optimizer should have two groups")
            self.assertEqual(len(optimizer_params_groups[0]["params"]), net.num_no_decay_params(),
                             msg="Wrong number of params for no decay group")
            self.assertEqual(len(optimizer_params_groups[1]["params"]), net.num_decay_params(),
                             msg="Wrong number of params for no decay group")
            self.assertEqual(optimizer_params_groups[0]['weight_decay'], 0,
                             msg="Weight decay value for no decay group should be zero")
            self.assertEqual(optimizer_params_groups[1]['weight_decay'], self.weight_decay,
                             msg="Wrong weight decay value for decay group.")

    def test_zero_wd_multiple_group(self):
        """
        test that 2 groups of parameters are separated to 2 groups of weight_decay_params and 2 groups without.
        """
        for input_dim in self.input_dimensions:
            net = ToySgModule(input_dimension=input_dim, multiple_param_groups=True)
            train_params = HpmStruct(**self.train_params_zero_wd)

            optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer(
                net,
                net.initialize_param_groups(self.lr, train_params),
                self.weight_decay
            )

            self.assertEqual(len(optimizer_params_groups), 4, msg="Optimizer should have 4 groups")
            self.assertEqual(len(optimizer_params_groups[0]["params"]), net.base_params[0],
                             msg="Wrong number of params for base no decay group")
            self.assertEqual(len(optimizer_params_groups[1]["params"]), net.base_params[1],
                             msg="Wrong number of params for base decay group")
            self.assertEqual(len(optimizer_params_groups[2]["params"]), net.head_params[0],
                             msg="Wrong number of params for head no decay group")
            self.assertEqual(len(optimizer_params_groups[3]["params"]), net.head_params[1],
                             msg="Wrong number of params for head decay group")
            self.assertEqual(optimizer_params_groups[0]['weight_decay'], 0,
                             msg="Weight decay value for no decay group should be zero")
            self.assertEqual(optimizer_params_groups[2]['weight_decay'], 0,
                             msg="Weight decay value for no decay group should be zero")
            # verifying that other parameters has a weight decay as intended
            self.assertEqual(optimizer_params_groups[1]['weight_decay'], self.weight_decay,
                             msg="Wrong weight decay value for decay group.")
            self.assertEqual(optimizer_params_groups[3]['weight_decay'], self.weight_decay,
                             msg="Wrong weight decay value for decay group.")


if __name__ == '__main__':
    unittest.main()