Problem with custom weight initialization with `SNES`

[FrizzoDavide]

Details:

• Library version: evotorch==0.5.1
• PyTorch version: torch==2.4.1
• Python version: 3.11.9
• Operating system: Ubuntu 24.04.1

Description:

I am having some issues with the creation of a custom searcher. I am working with SNES for Neuroevolution: I am doing Regression with time series data. In particular I would like to use a custom weight initialization so that I can use as initial candidate solutions the weight initialization of the nn.Module I am using instead of the initial population that is created sampling from the uniform hyper-cube provided by problem.initial_bounds.

So I created a custom class DamageSNES, which inherits SNES, where I only change the __init__ method. In particular to create a new initial population I did the following:

• First with self._population = self._problem.generate_batch(self._popsize) I created a SolutionBatch object (here we are however still using the default evotorch initialization based on initial_bounds)
• Then I created the initial weights I want to give to the different individuals of the population (i.e. the different neural networks) using the _initialize_model_weights method. This method is contained in a custom version of the SupervisedNE problem class (see below). For the moment _intialize_model_weights is still a dumb function that simply sets all the weights to 1 but I used it just to see if the initial weights changed.
• Then I exploited the set_values method to modify the values of self._population into the ones contained in initial_weights

Below I include the code I used, in particular:

• DamageSNES → Custom searcher class where I tried to implement custom weight initialization
• DamageSupervisedNE → Custom SupervisedNE class where I changed some of the basic methods of SupervisedNE and I added the _initialize_model_weights method to produce the initial weights.
• DamageLogger → This is a custom logger I use to log some results on wandb. Here I simply added a print statement at the beginning to test the custom weight initialization, printing at every iteration the current population

The problem I am facing is that it seems that the custom weight initialization is not happening and evotorch is still using the default initialization based on initial_bounds. In fact the DamageLogger is printing out populations composed of random value between the extremes of initial_bounds I provide in input and it is not printing a population full of 1 as it should happen considering how I built the _initialize_model_weights function.

class DamageSNES(SNES):
    def __init__(
        self,
        problem: DamageSupervisedNE,
        *,
        n_parameters: int,
        popsize: Optional[int] = None,
        radius_init: Optional[float] = None,
        center_learning_rate: Optional[float] = None,
        stdev_learning_rate: Optional[float] = None,
        scale_learning_rate: Optional[float] = None,
    ):
        super().__init__(
            problem=problem,
            popsize=popsize,
            radius_init=radius_init,
            center_learning_rate=center_learning_rate,
            stdev_learning_rate=stdev_learning_rate,
            scale_learning_rate=scale_learning_rate,
        )

        self._problem = problem
        self._popsize = popsize
        self._n_parameters = n_parameters

        self._population = self._problem.generate_batch(self._popsize)
        initial_weights = self._problem._initialize_model_weights(
            self._popsize, self._n_parameters
        )
        self._population.set_values(initial_weights)

class DamageSupervisedNE(SupervisedNE):
    def __init__(
        self,
        dataset: Dataset,
        network: nn.Module,
        loss_func: nn.Module,
        minibatch_size: int,
        num_minibatches: int,
        common_minibatch: bool,
        config: SimpleNamespace,
        device: str = "cpu",
        initial_bounds: Tuple[float, float] = (-0.05, 0.05),
        num_actors: int = 0,
    ):
        super().__init__(
            dataset=dataset,
            network=network,
            loss_func=loss_func,
            minibatch_size=minibatch_size,
            num_minibatches=num_minibatches,
            common_minibatch=common_minibatch,
            device=device,
            initial_bounds=initial_bounds,
            num_actors=num_actors,
        )

        self.dataset = dataset
        self.minibatch_size = minibatch_size
        self.config = config

    def _make_dataloader(self) -> DataLoader:
        return DataLoader(self.dataset, batch_size=self.minibatch_size, shuffle=False)

    def _evaluate_using_minibatch(
        self, network: nn.Module, batch: Any
    ) -> Union[float, torch.Tensor]:
        with torch.no_grad():
            x, y = batch
            if self.config.damage_loss:
                yhat, _ = network(x)
            elif self.config.signal_damage_loss:
                yhat, _, _ = network(x)
            else:
                yhat = network(x)

            return self.loss(yhat, y)

    def _initialize_model_weights(self, popsize: int, n_paramters: int) -> torch.Tensor:
        return torch.ones((popsize, n_paramters))

class DamageLogger(WandbLogger):
    def __init__(
        self,
        searcher,
        problem: DamageSupervisedNE,
        val_loader: DataLoader,
        test_loader: DataLoader,
        criterion: nn.Module,
        config: SimpleNamespace,
        **wandb_kwargs,
    ):
        super().__init__(searcher, **wandb_kwargs)
        self.searcher = searcher
        self.problem = problem
        self.val_loader = val_loader
        self.test_loader = test_loader
        self.criterion = criterion
        self.config = config
        self.best_state_dict: Dict = {}
        self.best_val_loss = math.inf

    def _log(
        self,
        status: Dict[str, Any],
    ):
        print("#" * 50)
        print("Current Population:")
        print(self.searcher._population.values)
        print("#" * 50)

        metrics_dict, state_dict, _, _ = evo_metrics(
            problem=self.problem,
            status_dict=self.searcher.status,
            val_loader=self.val_loader,
            test_loader=self.test_loader,
            criterion=self.criterion,
            config=self.config,
        )

        if metrics_dict["val_loss"] < self.best_val_loss:
            self.best_val_loss = metrics_dict["val_loss"]
            self.best_state_dict = state_dict

        status.update(metrics_dict)
        super()._log(status)

I don't know if custom weight initialization it's something that is possible to do in evotorch and weather there is a much easier way to do it than the approach I am trying to use.