(
        self,
        dataset: Dataset,
        epochs: List[int],
        batch_sizes: List[int],
        fade_in_percentages: List[int],
        loss_fn: GANLoss = WganGP(),
        batch_repeats: int = 4,
        gen_learning_rate: float = 0.003,
        dis_learning_rate: float = 0.003,
        num_samples: int = 16,
        start_depth: int = 2,
        num_workers: int = 3,
        feedback_factor: int = 100,
        save_dir=Path("./train"),
        checkpoint_factor: int = 10,
    )
# TODO implement support for conditional GAN here
        Utility method for training the ProGAN.
        Note that you don't have to necessarily use this method. You can use the
        optimize_generator and optimize_discriminator and define your own
        training routine
        Args:
            dataset: object of the dataset used for training.
                     Note that this is not the dataloader (we create dataloader in this
                     method since the batch_sizes for resolutions can be different)
            epochs: list of number of epochs to train the network for every resolution
            batch_sizes: list of batch_sizes for every resolution
            fade_in_percentages: list of percentages of epochs per resolution
                                used for fading in the new layer not used for
                                first resolution, but dummy value is still needed
            loss_fn: loss function used for training
            batch_repeats: number of iterations to perform on a single batch
            gen_learning_rate: generator learning rate
            dis_learning_rate: discriminator learning rate
            num_samples: number of samples generated in sample_sheet
            start_depth: start training from this depth
            num_workers: number of workers for reading the data
            feedback_factor: number of logs per epoch
            save_dir: directory for saving the models (.bin files)
            checkpoint_factor: save model after these many epochs.
        Returns: None (Writes multiple files to disk)
# TODO implement support for conditional GAN here
        Utility method for training the ProGAN.
        Note that you don't have to necessarily use this method. You can use the
        optimize_generator and optimize_discriminator and define your own
        training routine
        Args:
            dataset: object of the dataset used for training.
                     Note that this is not the dataloader (we create dataloader in this
                     method since the batch_sizes for resolutions can be different)
            epochs: list of number of epochs to train the network for every resolution
            batch_sizes: list of batch_sizes for every resolution
            fade_in_percentages: list of percentages of epochs per resolution
                                used for fading in the new layer not used for
                                first resolution, but dummy value is still needed
            loss_fn: loss function used for training
            batch_repeats: number of iterations to perform on a single batch
            gen_learning_rate: generator learning rate
            dis_learning_rate: discriminator learning rate
            num_samples: number of samples generated in sample_sheet
            start_depth: start training from this depth
            num_workers: number of workers for reading the data
            feedback_factor: number of logs per epoch
            save_dir: directory for saving the models (.bin files)
            checkpoint_factor: save model after these many epochs.
        Returns: None (Writes multiple files to disk)
# TODO implement support for conditional GAN here Utility method for training the ProGAN . Note that you don t have to necessarily use this method . You can use the optimize_generator and optimize_discriminator and define your own training routine Args : dataset : object of the dataset used for training . Note that this is not the dataloader ( we create dataloader in this method since the batch_sizes for resolutions can be different ) epochs : list of number of epochs to train the network for every resolution batch_sizes : list of batch_sizes for every resolution fade_in_percentages : list of percentages of epochs per resolution used for fading in the new layer not used for first resolution but dummy value is still needed loss_fn : loss function used for training batch_repeats : number of iterations to perform on a single batch gen_learning_rate : generator learning rate dis_learning_rate : discriminator learning rate num_samples : number of samples generated in sample_sheet start_depth : start training from this depth num_workers : number of workers for reading the data feedback_factor : number of logs per epoch save_dir : directory for saving the models ( . bin files ) checkpoint_factor : save model after these many epochs . Returns : None ( Writes multiple files to disk )
def train(
        self,
        dataset: Dataset,
        epochs: List[int],
        batch_sizes: List[int],
        fade_in_percentages: List[int],
        loss_fn: GANLoss = WganGP(),
        batch_repeats: int = 4,
        gen_learning_rate: float = 0.003,
        dis_learning_rate: float = 0.003,
        num_samples: int = 16,
        start_depth: int = 2,
        num_workers: int = 3,
        feedback_factor: int = 100,
        save_dir=Path("./train"),
        checkpoint_factor: int = 10,
    ):
        """
        # TODO implement support for conditional GAN here
        Utility method for training the ProGAN.
        Note that you don't have to necessarily use this method. You can use the
        optimize_generator and optimize_discriminator and define your own
        training routine
        Args:
            dataset: object of the dataset used for training.
                     Note that this is not the dataloader (we create dataloader in this
                     method since the batch_sizes for resolutions can be different)
            epochs: list of number of epochs to train the network for every resolution
            batch_sizes: list of batch_sizes for every resolution
            fade_in_percentages: list of percentages of epochs per resolution
                                used for fading in the new layer not used for
                                first resolution, but dummy value is still needed
            loss_fn: loss function used for training
            batch_repeats: number of iterations to perform on a single batch
            gen_learning_rate: generator learning rate
            dis_learning_rate: discriminator learning rate
            num_samples: number of samples generated in sample_sheet
            start_depth: start training from this depth
            num_workers: number of workers for reading the data
            feedback_factor: number of logs per epoch
            save_dir: directory for saving the models (.bin files)
            checkpoint_factor: save model after these many epochs.
        Returns: None (Writes multiple files to disk)
        """

        print(f"Loaded the dataset with: {len(dataset)} images ...")
        assert (self.depth - 1) == len(
            batch_sizes
        ), "batch_sizes are not compatible with depth"
        assert (self.depth - 1) == len(epochs), "epochs are not compatible with depth"

        self._toggle_all_networks("train")

        # create the generator and discriminator optimizers
        gen_optim = torch.optim.Adam(
            params=self.gen.parameters(),
            lr=gen_learning_rate,
            betas=(0, 0.99),
            eps=1e-8,
        )
        dis_optim = torch.optim.Adam(
            params=self.dis.parameters(),
            lr=dis_learning_rate,
            betas=(0, 0.99),
            eps=1e-8,
        )

        # verbose stuff
        print("setting up the image saving mechanism")
        model_dir, log_dir = save_dir / "models", save_dir / "logs"
        model_dir.mkdir(parents=True, exist_ok=True)
        log_dir.mkdir(parents=True, exist_ok=True)

        feedback_generator = self.gen_shadow if self.use_ema else self.gen

        # image saving mechanism
        with torch.no_grad():
            dummy_data_loader = get_data_loader(dataset, num_samples, num_workers)
            real_images_for_render = next(iter(dummy_data_loader))
            fixed_input = torch.randn(num_samples, self.latent_size).to(self.device)
            self.create_grid(
                real_images_for_render,
                scale_factor=1,
                img_file=log_dir / "real_images.png",
            )
            self.create_grid(
                feedback_generator(fixed_input, self.depth, 1).detach(),
                scale_factor=1,
                img_file=log_dir / "fake_images_0.png",
            )

        # tensorboard summarywriter:
        summary = SummaryWriter(str(log_dir / "tensorboard"))

        # create a global time counter
        global_time = time.time()
        global_step = 0

        print("Starting the training process ... ")
        for current_depth in range(start_depth, self.depth + 1):
            current_res = int(2 ** current_depth)
            print(f"\n\nCurrently working on Depth: {current_depth}")
            print("Current resolution: %d x %d" % (current_res, current_res))
            depth_list_index = current_depth - 2
            current_batch_size = batch_sizes[depth_list_index]
            data = get_data_loader(dataset, current_batch_size, num_workers)
            ticker = 1
            for epoch in range(1, epochs[depth_list_index] + 1):
                start = timeit.default_timer()  # record time at the start of epoch
                print(f"\nEpoch: {epoch}")
                total_batches = len(data)

                # compute the fader point
                fader_point = int(
                    (fade_in_percentages[depth_list_index] / 100)
                    * epochs[depth_list_index]
                    * total_batches
                )

                for (i, batch) in enumerate(data, start=1):
                    # calculate the alpha for fading in the layers
                    alpha = ticker / fader_point if ticker <= fader_point else 1

                    # extract current batch of data for training
                    images = batch.to(self.device)

                    gan_input = torch.randn(current_batch_size, self.latent_size).to(
                        self.device
                    )

                    gen_loss, dis_loss = None, None
                    for _ in range(batch_repeats):
                        # optimize the discriminator:
                        dis_loss = self.optimize_discriminator(
                            loss_fn, dis_optim, gan_input, images, current_depth, alpha
                        )

                        # no idea why this needs to be done after discriminator
                        # iteration, but this is where it is done in the original
                        # code
                        if self.use_ema:
                            update_average(
                                self.gen_shadow, self.gen, beta=self.ema_beta
                            )

                        # optimize the generator:
                        gen_loss = self.optimize_generator(
                            loss_fn, gen_optim, gan_input, images, current_depth, alpha
                        )
                    global_step += 1

                    # provide a loss feedback
                    if (
                        i % max(int(total_batches / max(feedback_factor, 1)), 1) == 0
                        or i == 1
                        or i == total_batches
                    ):
                        elapsed = time.time() - global_time
                        elapsed = str(datetime.timedelta(seconds=elapsed))
                        print(
                            "Elapsed: [%s]  batch: %d  d_loss: %f  g_loss: %f"
                            % (elapsed, i, dis_loss, gen_loss)
                        )
                        summary.add_scalar(
                            "dis_loss", dis_loss, global_step=global_step
                        )
                        summary.add_scalar(
                            "gen_loss", gen_loss, global_step=global_step
                        )
                        # create a grid of samples and save it
                        resolution_dir = log_dir / str(int(2 ** current_depth))
                        resolution_dir.mkdir(exist_ok=True)
                        gen_img_file = resolution_dir / f"{epoch}_{i}.png"

                        # this is done to allow for more GPU space
                        with torch.no_grad():
                            self.create_grid(
                                samples=feedback_generator(
                                    fixed_input, current_depth, alpha
                                ).detach(),
                                scale_factor=int(2 ** (self.depth - current_depth)),
                                img_file=gen_img_file,
                            )

                    # increment the alpha ticker and the step
                    ticker += 1

                stop = timeit.default_timer()
                print("Time taken for epoch: %.3f secs" % (stop - start))

                if (
                    epoch % checkpoint_factor == 0
                    or epoch == 1
                    or epoch == epochs[depth_list_index]
                ):
                    save_file = model_dir / f"depth_{current_depth}_epoch_{epoch}.bin"
                    torch.save(self.get_save_info(gen_optim, dis_optim), save_file)

        self._toggle_all_networks("eval")
        print("Training completed ...")
def train ( self , dataset : Dataset , epochs : List [ int ] , batch_sizes : List [ int ] , fade_in_percentages : List [ int ] , loss_fn : GANLoss = WganGP ( ) , batch_repeats : int = 4 , gen_learning_rate : float = 0.003 , dis_learning_rate : float = 0.003 , num_samples : int = 16 , start_depth : int = 2 , num_workers : int = 3 , feedback_factor : int = 100 , save_dir = Path ( "./train" ) , checkpoint_factor : int = 10 , ) : print ( f"Loaded the dataset with: {len(dataset)} images ..." ) assert ( self . depth - 1 ) == len ( batch_sizes ) , "batch_sizes are not compatible with depth" assert ( self . depth - 1 ) == len ( epochs ) , "epochs are not compatible with depth" self . _toggle_all_networks ( "train" ) # create the generator and discriminator optimizers gen_optim = torch . optim . Adam ( params = self . gen . parameters ( ) , lr = gen_learning_rate , betas = ( 0 , 0.99 ) , eps = 1e-8 , ) dis_optim = torch . optim . Adam ( params = self . dis . parameters ( ) , lr = dis_learning_rate , betas = ( 0 , 0.99 ) , eps = 1e-8 , ) # verbose stuff print ( "setting up the image saving mechanism" ) model_dir , log_dir = save_dir / "models" , save_dir / "logs" model_dir . mkdir ( parents = True , exist_ok = True ) log_dir . mkdir ( parents = True , exist_ok = True ) feedback_generator = self . gen_shadow if self . use_ema else self . gen # image saving mechanism with torch . no_grad ( ) : dummy_data_loader = get_data_loader ( dataset , num_samples , num_workers ) real_images_for_render = next ( iter ( dummy_data_loader ) ) fixed_input = torch . randn ( num_samples , self . latent_size ) . to ( self . device ) self . create_grid ( real_images_for_render , scale_factor = 1 , img_file = log_dir / "real_images.png" , ) self . create_grid ( feedback_generator ( fixed_input , self . depth , 1 ) . detach ( ) , scale_factor = 1 , img_file = log_dir / "fake_images_0.png" , ) # tensorboard summarywriter: summary = SummaryWriter ( str ( log_dir / "tensorboard" ) ) # create a global time counter global_time = time . time ( ) global_step = 0 print ( "Starting the training process ... " ) for current_depth in range ( start_depth , self . depth + 1 ) : current_res = int ( 2 ** current_depth ) print ( f"\n\nCurrently working on Depth: {current_depth}" ) print ( "Current resolution: %d x %d" % ( current_res , current_res ) ) depth_list_index = current_depth - 2 current_batch_size = batch_sizes [ depth_list_index ] data = get_data_loader ( dataset , current_batch_size , num_workers ) ticker = 1 for epoch in range ( 1 , epochs [ depth_list_index ] + 1 ) : start = timeit . default_timer ( ) # record time at the start of epoch print ( f"\nEpoch: {epoch}" ) total_batches = len ( data ) # compute the fader point fader_point = int ( ( fade_in_percentages [ depth_list_index ] / 100 ) * epochs [ depth_list_index ] * total_batches ) for ( i , batch ) in enumerate ( data , start = 1 ) : # calculate the alpha for fading in the layers alpha = ticker / fader_point if ticker <= fader_point else 1 # extract current batch of data for training images = batch . to ( self . device ) gan_input = torch . randn ( current_batch_size , self . latent_size ) . to ( self . device ) gen_loss , dis_loss = None , None for _ in range ( batch_repeats ) : # optimize the discriminator: dis_loss = self . optimize_discriminator ( loss_fn , dis_optim , gan_input , images , current_depth , alpha ) # no idea why this needs to be done after discriminator # iteration, but this is where it is done in the original # code if self . use_ema : update_average ( self . gen_shadow , self . gen , beta = self . ema_beta ) # optimize the generator: gen_loss = self . optimize_generator ( loss_fn , gen_optim , gan_input , images , current_depth , alpha ) global_step += 1 # provide a loss feedback if ( i % max ( int ( total_batches / max ( feedback_factor , 1 ) ) , 1 ) == 0 or i == 1 or i == total_batches ) : elapsed = time . time ( ) - global_time elapsed = str ( datetime . timedelta ( seconds = elapsed ) ) print ( "Elapsed: [%s]  batch: %d  d_loss: %f  g_loss: %f" % ( elapsed , i , dis_loss , gen_loss ) ) summary . add_scalar ( "dis_loss" , dis_loss , global_step = global_step ) summary . add_scalar ( "gen_loss" , gen_loss , global_step = global_step ) # create a grid of samples and save it resolution_dir = log_dir / str ( int ( 2 ** current_depth ) ) resolution_dir . mkdir ( exist_ok = True ) gen_img_file = resolution_dir / f"{epoch}_{i}.png" # this is done to allow for more GPU space with torch . no_grad ( ) : self . create_grid ( samples = feedback_generator ( fixed_input , current_depth , alpha ) . detach ( ) , scale_factor = int ( 2 ** ( self . depth - current_depth ) ) , img_file = gen_img_file , ) # increment the alpha ticker and the step ticker += 1 stop = timeit . default_timer ( ) print ( "Time taken for epoch: %.3f secs" % ( stop - start ) ) if ( epoch % checkpoint_factor == 0 or epoch == 1 or epoch == epochs [ depth_list_index ] ) : save_file = model_dir / f"depth_{current_depth}_epoch_{epoch}.bin" torch . save ( self . get_save_info ( gen_optim , dis_optim ) , save_file ) self . _toggle_all_networks ( "eval" ) print ( "Training completed ..." )