def __init__(self, logf0s_normalization, mcep_normalization, model_checkpoint, validation_A_dir, output_A_dir): logf0s_normalization = np.load(logf0s_normalization) self.log_f0s_mean_A = logf0s_normalization['mean_A'] self.log_f0s_std_A = logf0s_normalization['std_A'] self.log_f0s_mean_B = logf0s_normalization['mean_B'] self.log_f0s_std_B = logf0s_normalization['std_B'] mcep_normalization = np.load(mcep_normalization) self.coded_sps_A_mean = mcep_normalization['mean_A'] self.coded_sps_A_std = mcep_normalization['std_A'] self.coded_sps_B_mean = mcep_normalization['mean_B'] self.coded_sps_B_std = mcep_normalization['std_B'] self.validation_A_dir = validation_A_dir self.output_A_dir = output_A_dir self.device = torch.device( 'cuda' if torch.cuda.is_available() else 'cpu') self.generator_A2B = Generator().to(self.device) self.generator_A2B.eval() checkPoint = torch.load(model_checkpoint) self.generator_A2B.load_state_dict( state_dict=checkPoint['model_genA2B_state_dict']) print("load model checkpoint finish!")
def __init__(self, logf0s_normalization, mcep_normalization, coded_sps_A_norm, coded_sps_B_norm, model_checkpoint, validation_A_dir, output_A_dir, validation_B_dir, output_B_dir, restart_training_at):#=None): self.start_epoch = 0 self.num_epochs = 2000 self.mini_batch_size = 16 self.dataset_A = self.loadPickleFile(coded_sps_A_norm) self.dataset_B = self.loadPickleFile(coded_sps_B_norm) self.device = torch.device( 'cuda' if torch.cuda.is_available() else 'cpu') # Speech Parameters logf0s_normalization = np.load(logf0s_normalization) self.log_f0s_mean_A = logf0s_normalization['mean_A'] self.log_f0s_std_A = logf0s_normalization['std_A'] self.log_f0s_mean_B = logf0s_normalization['mean_B'] self.log_f0s_std_B = logf0s_normalization['std_B'] mcep_normalization = np.load(mcep_normalization) self.coded_sps_A_mean = mcep_normalization['mean_A'] self.coded_sps_A_std = mcep_normalization['std_A'] self.coded_sps_B_mean = mcep_normalization['mean_B'] self.coded_sps_B_std = mcep_normalization['std_B'] # Generator and Discriminator self.generator_A2B = Generator().to(self.device) self.generator_B2A = Generator().to(self.device) self.discriminator_A = Discriminator().to(self.device) self.discriminator_B = Discriminator().to(self.device) # Loss Functions criterion_mse = torch.nn.MSELoss() # Optimizer g_params = list(self.generator_A2B.parameters()) + \ list(self.generator_B2A.parameters()) d_params = list(self.discriminator_A.parameters()) + \ list(self.discriminator_B.parameters()) # Initial learning rates self.generator_lr = 0.0002 self.discriminator_lr = 0.0001 # Learning rate decay self.generator_lr_decay = self.generator_lr / 200000 self.discriminator_lr_decay = self.discriminator_lr / 200000 # Starts learning rate decay from after this many iterations have passed self.start_decay = 12500 self.generator_optimizer = torch.optim.Adam( g_params, lr=self.generator_lr, betas=(0.5, 0.999)) self.discriminator_optimizer = torch.optim.Adam( d_params, lr=self.discriminator_lr, betas=(0.5, 0.999)) # To Load save previously saved models self.modelCheckpoint = model_checkpoint # Validation set Parameters self.validation_A_dir = validation_A_dir self.output_A_dir = output_A_dir self.validation_B_dir = validation_B_dir self.output_B_dir = output_B_dir # Storing Discriminatior and Generator Loss self.generator_loss_store = [] self.discriminator_loss_store = [] self.file_name = 'log_store_non_sigmoid.txt' if restart_training_at is not None: # Training will resume from previous checkpoint self.start_epoch = self.loadModel(restart_training_at) print("Training resumed")