def init_train_set(epoch, from_iter):
#train_dataset.set_curriculum_epoch(epoch, sample=True)
train_dataset.set_curriculum_epoch(epoch, sample=False)
global train_loader, train_sampler
if not args.distributed:
train_sampler = BucketingSampler(train_dataset,
batch_size=args.batch_size)
train_sampler.bins = train_sampler.bins[from_iter:]
else:
train_sampler = DistributedBucketingSampler(
train_dataset,
batch_size=args.batch_size,
num_replicas=args.world_size,
rank=args.rank)
train_loader = AudioDataLoader(train_dataset,
num_workers=args.num_workers,
batch_sampler=train_sampler)
if (not args.no_shuffle and epoch != 0) or args.no_sorta_grad:
print("Shuffling batches for the following epochs")
train_sampler.shuffle(epoch)
for g in param_groups:
g['lr'] = g['lr'] / args.learning_anneal
if g['lr'] <= 0.0001:
g['lr'] = 0.0001
print('Learning rate annealed to: {lr:.6f}'.format(lr=g['lr']))
if main_proc and (best_wer is None or best_wer > wer):
print("Found better validated model, saving to %s" %
args.model_path)
#torch.save(DeepSpeech.serialize(model, optimizer=optimizer, epoch=epoch, loss_results=loss_results,
# wer_results=wer_results, cer_results=cer_results)
# , args.model_path)
torch.save(denoiser.state_dict(), args.save_path)
best_wer = wer
avg_loss = 0
if not args.no_shuffle:
print("Shuffling batches...")
g = torch.Generator()
g.manual_seed(epoch)
bins_ids = list(
torch.randperm(len(train_sampler_adv.bins), generator=g))
#bins_ids.append(torch.tensor(len(train_sampler_noise.bins)-1))
train_sampler_clean.bins = [
train_sampler_clean.bins[x] for x in bins_ids
]
train_sampler_adv.bins = [
train_sampler_adv.bins[x] for x in bins_ids
]