if args.sample_strategy == 'sequential':
sample_strategy = sampler.SequentialSampler(dataset)
elif args.sample_strategy == 'random':
sample_strategy = sampler.RandomSampler(dataset)
dataloader = DataLoader(dataset,
batch_size=params['batch_size'],
num_workers=params['num_workers'],
sampler=sample_strategy)
dummy_input, _, _, _, _, dummy_one_hot = dataset[0]
params['num_attractors'] = dummy_one_hot.shape[-1]
params['num_sources'] = params['num_attractors']
params['sample_rate'] = dataset.sr
dataset.reorder_sources = args.reorder_sources
val_dataset.reorder_sources = args.reorder_sources
pp.pprint(params)
class_func = MaskEstimation if args.baseline else DeepAttractor
model = utils.load_class_from_params(params, class_func).to(device)
if not params['compute_training_stats']:
mean, std = utils.compute_statistics_on_dataset(dataloader, device)
params['training_stats'] = {'mean': mean, 'std': std + 1e-7}
dataset.stats = params['training_stats']
val_dataset.stats = params['training_stats']
dataset.whiten_data = True
val_dataset.whiten_data = True