def __init__(self, train_loader, test_loader, config, load_dir=None):
self.train_loader = train_loader
self.test_loader = test_loader
self.sample_set = Sample_Set(config)
self.config = config
self.model_name = self.config['model']['name']
self.model = model.StarGAN_emo_VC1(self.config, self.model_name)
self.set_configuration()
self.model = self.model
if not load_dir == None:
self.load_checkpoint(load_dir)
np.random.seed(SEED)
random.seed(SEED)
# Use GPU
USE_GPU = True
if USE_GPU and torch.cuda.is_available():
device = torch.device('cuda')
torch.cuda.manual_seed_all(SEED)
map_location='cuda'
else:
device = torch.device('cpu')
map_location='cpu'
# Load model
model = model.StarGAN_emo_VC1(config, config['model']['name'])
# model.load(args.checkpoint)
model.load(checkpoint_dir, map_location= map_location)
config = model.config
model.to_device(device = device)
model.set_eval_mode()
# Make emotion targets (using config file)
# s = solver.Solver(None, None, config, load_dir = None)
# targets =
num_emos = config['model']['num_classes']
emo_labels = torch.Tensor(range(0,num_emos)).long()
emo_targets = F.one_hot(emo_labels, num_classes = num_emos).float().to(device = device)
print(f"Number of emotions = {num_emos}")
if args.in_dir == 'sample':