Method reset

Image_Classification/src/train.py:202–246

(self, learning_rate: float)

            raise NotImplementedError(f"{self.config['network']} does not support pretrained weights yet")

    def reset(self, learning_rate: float) -> None:
        self.network = get_network(name=self.config['network'], num_classes=self.num_classes)
        for layer in self.network.modules():
            if isinstance(layer, nn.BatchNorm2d):
                layer.float()
        if self.pretrained:
            self.get_pretrained_model()
        if self.device == 'cpu':
            print("Resetting cpu-based network")
        elif self.device == 'mps':
            self.network = self.network.to(self.gpu)
        elif self.dist:
            self.network.to(self.gpu)
            self.network = torch.nn.parallel.DistributedDataParallel(
                self.network, device_ids=[self.gpu])
        else:
            self.network = self.network.cuda(self.gpu)

        self.optimizer, self.scheduler = get_optimizer_scheduler(
            optim_method=self.config['optimizer'],
            lr_scheduler=self.config['scheduler'],
            init_lr=learning_rate,
            net=self.network,
            train_loader_len=len(self.train_loader),
            max_epochs=self.config['max_epochs'],
            optimizer_kwargs=self.config['optimizer_kwargs'],
            scheduler_kwargs=self.config['scheduler_kwargs'])
        if self.config['optimizer'] in ['Shampoo', 'kfac']:
            damping = self.config['optimizer_kwargs']['damping']
            curvature_update_interval = self.config['optimizer_kwargs']['curvature_update_interval']
            ema_decay = self.config['optimizer_kwargs']['ema_decay']
            config = PreconditioningConfig(data_size=self.config['mini_batch_size'],
                                           damping=damping,
                                           preconditioner_upd_interval=curvature_update_interval,
                                           curvature_upd_interval=curvature_update_interval,
                                           ema_decay=ema_decay,
                                           ignore_modules=[nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d,
                                                           nn.LayerNorm])
            if self.config['optimizer'] == "Shampoo":
                self.gm = ShampooGradientMaker(self.network, config)
            elif self.config['optimizer'] == "kfac":
                self.gm = KfacGradientMaker(self.network, config)
            else:
                raise ValueError(f"This optimizer is not reconized: {self.config['optimizer']}")
        self.early_stop.reset()

    def create_output_dir(self, checkpoint: bool = False):
        def create_dir(path: Path):