def train_one_epoch(self, criterion, data_loaders, phases=('train', 'test')):
epoch_beg = time.time()
self.model.to(self.device)
dataset_sizes = {x: len(data_loaders[x].dataset)
for x in phases}
batch_sizes = {x: data_loaders[x].batch_size for x in phases}
try:
m = (1 - self.params_dict['warmup_initial_scale'])/(self.params_dict['warmup_num_epochs'] - 1)
c = self.params_dict['warmup_initial_scale']
except ZeroDivisionError:
m = 1
c = 0
lr_scale = m * self.elapsed_epochs() + c
if self.elapsed_epochs() < self.params_dict['warmup_num_epochs']:
self.optimizer.param_groups[0]['lr'] *= lr_scale
self.loss_history[str(self.elapsed_epochs())] = {}
print('Epoch {}/{} - lr={}'.format(self.elapsed_epochs(), self.params_dict['num_epochs'],
self.optimizer.param_groups[0]['lr'])
)
for phase in phases:
print('\t{} '.format(phase.title()), end='')
phase_beg = time.time()
if phase == 'train':
self.model.train()
else:
self.model.eval()
running_loss = 0.
running_dist_sum = 0.
running_count = 0.
running_count_per_class = torch.zeros(self.model.num_centroids)
for batch_no, batch_data in enumerate(data_loaders[phase]):
self.optimizer.zero_grad()
data_batch, label_batch, speaker_ids = batch_data
data_batch = data_batch.to(self.device)
label_batch = label_batch.to(self.device)
speaker_ids = speaker_ids.to(self.device)
with torch.set_grad_enabled(phase == 'train'):
encoded = self.model.encoder(data_batch).contiguous()
quantized, alignment = self.model.quantize(encoded, return_alignment=True)
count_per_class = torch.tensor([(alignment == i).sum()
for i in range(self.model.num_centroids)])
running_count_per_class += count_per_class
output = self.model.decoder(quantized)
predicted_centroids = self.model.centroids[alignment]
encoded = encoded.view(-1, encoded.size(-1))
new_length = output.size(1)
old_length = label_batch.size(1)
# Try to correct slight mismatches from down-sampling and up-sampling.
# Be careful in case of down-sampling without up-sampling
if old_length > new_length:
label_batch = label_batch[:, :new_length]
elif new_length > old_length:
output = output[:, :old_length]
distance_loss = criterion['dis_loss'](encoded, predicted_centroids.detach())
commitment_loss = criterion['com_loss'](predicted_centroids, encoded.detach())
reconstruction_loss = criterion['rec_loss'](output, label_batch)
if self.model.use_ma:
total_loss = reconstruction_loss + distance_loss
else:
total_loss = reconstruction_loss + distance_loss + commitment_loss
# Not part of the graph, just metrics to be monitored
mask = (label_batch != data_loaders[phase].dataset.pad_value).float()
numel = 1 # torch.sum(mask).item()
running_loss += reconstruction_loss.item()
running_dist_sum += distance_loss.item()
running_count += numel
class0_sum = running_dist_sum / running_count
px = running_count_per_class / running_count_per_class.sum()
px.clamp_min_(1e-20)
entropy = -(px * torch.log2(px)).sum().item()
if phase == 'train':
total_loss.backward()
if self.params_dict['gradient_clip']:
torch.nn.utils.clip_grad_norm_(self.model.parameters(),
self.params_dict['gradient_clip'])
self.optimizer.step()
phase_elapse = since(phase_beg)
eta = int(phase_elapse
* (dataset_sizes[phase] // batch_sizes[phase]
- batch_no - 1)
/ (batch_no + 1))
if self.params_dict['verbose']:
print('\r\t{} batch: {}/{} batches - ETA: {}s - loss: {:.4f} - dist: {:.4f} - '
'entropy: {:.4f}'.format(phase.title(),
batch_no + 1,
dataset_sizes[phase] // batch_sizes[phase] + 1,
eta, running_loss/running_count,
class0_sum, entropy
), end='')
epoch_loss = running_loss/running_count
print(" - loss: {:.4f} - dist: {:.4f} - entropy: {:.4f}".format(epoch_loss, class0_sum,
entropy))
self.loss_history[str(self.elapsed_epochs() - 1)][phase] = (epoch_loss,
class0_sum, entropy,
running_count)
print('\tTime: {}s'.format(int(since(epoch_beg))))
if self.elapsed_epochs() <= self.params_dict['warmup_num_epochs']:
self.optimizer.param_groups[0]['lr'] /= lr_scale
def test_since_second():
ref = datetime.datetime(2019, 5, 16, 13, 35, 15)
now = datetime.datetime(2019, 5, 16, 13, 35, 16)
assert utils.since(now, ref) == '1 second'
def test_since_hour_and_minutes_and_seconds():
ref = datetime.datetime(2019, 5, 16, 13, 35, 15)
now = datetime.datetime(2019, 5, 16, 14, 37, 21)
assert utils.since(now, ref) == '1 hour 2 minutes 6 seconds'
def test_since_days_hours_and_minutes_and_seconds():
ref = datetime.datetime(2019, 5, 16, 13, 35, 15)
now = datetime.datetime(2019, 5, 19, 15, 37, 21)
assert utils.since(now, ref) == '3 days 2 hours 2 minutes 6 seconds'
def test_since_minute_and_seconds():
ref = datetime.datetime(2019, 5, 16, 13, 35, 15)
now = datetime.datetime(2019, 5, 16, 13, 36, 21)
assert utils.since(now, ref) == '1 minute 6 seconds'
