def get_preprocessor(self, online_config):
# load the same preprocessor as pretraining stage
upstream_input_feat = online_config['input']
upstream_input_feat['channel'] = 0
upstream_target_feat = online_config['target']
upstream_target_feat['channel'] = 0
preprocessor = OnlinePreprocessor(**online_config, feat_list=[upstream_input_feat, upstream_target_feat])
upstream_feat = preprocessor()[0]
upstream_input_dim = upstream_feat.size(-1)
return preprocessor, upstream_input_dim
def __init__(self,
file_path,
sets,
bucket_size,
max_timestep=0,
drop=False,
mam_config=None,
libri_root=None,
online_config=None):
super(AcousticDataset, self).__init__(file_path, sets, bucket_size,
max_timestep, drop)
self.mam_config = mam_config
self.libri_root = libri_root
self.online_config = online_config
if self.online_config is not None:
assert libri_root is not None
feat_list = [
self.online_config['input'], self.online_config['target']
]
self.preprocessor = OnlinePreprocessor(**self.online_config,
feat_list=feat_list)
self.sample_step = mam_config[
'max_input_length'] if 'max_input_length' in mam_config else 0
if self.sample_step > 0:
print(
'[Dataset] - Sampling random segments for training, sample length:',
self.sample_step)
X = self.table['file_path'].tolist()
X_lens = self.table['length'].tolist()
# Use bucketing to allow different batch size at run time
self.X = []
batch_x, batch_len = [], []
for x, x_len in zip(X, X_lens):
batch_x.append(x)
batch_len.append(x_len)
# Fill in batch_x until batch is full
if len(batch_x) == bucket_size:
# Half the batch size if seq too long
if (bucket_size >= 2) and (max(batch_len) > HALF_BATCHSIZE_TIME
) and self.sample_step == 0:
self.X.append(batch_x[:bucket_size // 2])
self.X.append(batch_x[bucket_size // 2:])
else:
self.X.append(batch_x)
batch_x, batch_len = [], []
# Gather the last batch
if len(batch_x) > 1:
self.X.append(batch_x)
def __init__(self, args, config, dataloader, ckpdir):
self.device = torch.device('cuda') if (args.gpu and torch.cuda.is_available()) else torch.device('cpu')
if torch.cuda.is_available(): print('[Runner] - CUDA is available!')
self.model_kept = []
self.global_step = 1
self.log = SummaryWriter(ckpdir)
self.args = args
self.config = config
self.dataloader = dataloader
self.ckpdir = ckpdir
# optimizer
self.learning_rate = float(config['optimizer']['learning_rate'])
self.warmup_proportion = config['optimizer']['warmup_proportion']
self.gradient_accumulation_steps = config['optimizer']['gradient_accumulation_steps']
self.gradient_clipping = config['optimizer']['gradient_clipping']
# Training details
self.apex = config['runner']['apex']
self.total_steps = config['runner']['total_steps']
self.log_step = config['runner']['log_step']
self.save_step = config['runner']['save_step']
self.duo_feature = config['runner']['duo_feature']
self.max_keep = config['runner']['max_keep']
# model
self.transformer_config = config['transformer']
if args.online_feat is not None:
feat_list = [args.online_feat['input'], args.online_feat['target']]
self.preprocessor = OnlinePreprocessor(**args.online_feat, feat_list=feat_list).to(device=self.device)
self.input_dim, self.output_dim = [feat.size(-1) for feat in self.preprocessor()]
self.preprocessor.test_istft()
else:
self.input_dim = self.transformer_config['input_dim']
self.output_dim = 1025 if self.duo_feature else None # output dim is the same as input dim if not using duo features
class Runner():
''' Handler for complete pre-training progress of upstream models '''
def __init__(self, args, config, dataloader, ckpdir):
self.device = torch.device('cuda') if (args.gpu and torch.cuda.is_available()) else torch.device('cpu')
if torch.cuda.is_available(): print('[Runner] - CUDA is available!')
self.model_kept = []
self.global_step = 1
self.log = SummaryWriter(ckpdir)
self.args = args
self.config = config
self.dataloader = dataloader
self.ckpdir = ckpdir
# optimizer
self.learning_rate = float(config['optimizer']['learning_rate'])
self.warmup_proportion = config['optimizer']['warmup_proportion']
self.gradient_accumulation_steps = config['optimizer']['gradient_accumulation_steps']
self.gradient_clipping = config['optimizer']['gradient_clipping']
# Training details
self.apex = config['runner']['apex']
self.total_steps = config['runner']['total_steps']
self.log_step = config['runner']['log_step']
self.save_step = config['runner']['save_step']
self.duo_feature = config['runner']['duo_feature']
self.max_keep = config['runner']['max_keep']
# model
self.transformer_config = config['transformer']
if args.online_feat is not None:
feat_list = [args.online_feat['input'], args.online_feat['target']]
self.preprocessor = OnlinePreprocessor(**args.online_feat, feat_list=feat_list).to(device=self.device)
self.input_dim, self.output_dim = [feat.size(-1) for feat in self.preprocessor()]
self.preprocessor.test_istft()
else:
self.input_dim = self.transformer_config['input_dim']
self.output_dim = 1025 if self.duo_feature else None # output dim is the same as input dim if not using duo features
def set_model(self):
print('[Runner] - Initializing Transformer model...')
# build the Transformer model with speech prediction head
model_config = TransformerConfig(self.config)
self.dr = model_config.downsample_rate
self.model = TransformerForMaskedAcousticModel(model_config, self.input_dim, self.output_dim).to(self.device)
self.model.train()
if self.args.multi_gpu:
self.model = torch.nn.DataParallel(self.model)
print('[Runner] - Multi-GPU training Enabled: ' + str(torch.cuda.device_count()))
print('[Runner] - Number of parameters: ' + str(sum(p.numel() for p in self.model.parameters() if p.requires_grad)))
# Setup optimizer
param_optimizer = list(self.model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
if self.apex:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=self.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if self.config['optimizer']['loss_scale'] == 0:
self.optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
self.optimizer = FP16_Optimizer(optimizer, static_loss_scale=self.config['optimizer']['loss_scale'])
self.warmup_linear = WarmupLinearSchedule(warmup=self.warmup_proportion,
t_total=self.total_steps)
else:
self.optimizer = BertAdam(optimizer_grouped_parameters,
lr=self.learning_rate,
warmup=self.warmup_proportion,
t_total=self.total_steps)
def save_model(self, name='states', to_path=None):
all_states = {
'SpecHead': self.model.SpecHead.state_dict() if not self.args.multi_gpu else self.model.module.SpecHead.state_dict(),
'Transformer': self.model.Transformer.state_dict() if not self.args.multi_gpu else self.model.module.Transformer.state_dict(),
'Optimizer': self.optimizer.state_dict(),
'Global_step': self.global_step,
'Settings': {
'Config': self.config,
'Paras': self.args,
},
}
if to_path is None:
new_model_path = '{}/{}-{}.ckpt'.format(self.ckpdir, name, self.global_step)
else:
new_model_path = to_path
torch.save(all_states, new_model_path)
self.model_kept.append(new_model_path)
if len(self.model_kept) >= self.max_keep:
os.remove(self.model_kept[0])
self.model_kept.pop(0)
def up_sample_frames(self, spec, return_first=False):
if len(spec.shape) != 3:
spec = spec.unsqueeze(0)
assert(len(spec.shape) == 3), 'Input should have acoustic feature of shape BxTxD'
# spec shape: [batch_size, sequence_length // downsample_rate, output_dim * downsample_rate]
spec_flatten = spec.view(spec.shape[0], spec.shape[1]*self.dr, spec.shape[2]//self.dr)
if return_first: return spec_flatten[0]
return spec_flatten # spec_flatten shape: [batch_size, sequence_length * downsample_rate, output_dim // downsample_rate]
def down_sample_frames(self, spec):
left_over = spec.shape[1] % self.dr
if left_over != 0: spec = spec[:, :-left_over, :]
spec_stacked = spec.view(spec.shape[0], spec.shape[1]//self.dr, spec.shape[2]*self.dr)
return spec_stacked
def process_data(self, spec):
"""Process training data for the masked acoustic model"""
with torch.no_grad():
assert(len(spec) == 5), 'dataloader should return (spec_masked, pos_enc, mask_label, attn_mask, spec_stacked)'
# Unpack and Hack bucket: Bucketing should cause acoustic feature to have shape 1xBxTxD'
spec_masked = spec[0].squeeze(0)
pos_enc = spec[1].squeeze(0)
mask_label = spec[2].squeeze(0)
attn_mask = spec[3].squeeze(0)
spec_stacked = spec[4].squeeze(0)
spec_masked = spec_masked.to(device=self.device)
if pos_enc.dim() == 3:
# pos_enc: (batch_size, seq_len, hidden_size)
# GPU memory need (batch_size * seq_len * hidden_size)
pos_enc = pos_enc.float().to(device=self.device)
elif pos_enc.dim() == 2:
# pos_enc: (seq_len, hidden_size)
# GPU memory only need (seq_len * hidden_size) even after expanded
pos_enc = pos_enc.float().to(device=self.device).expand(spec_masked.size(0), *pos_enc.size())
mask_label = mask_label.bool().to(device=self.device)
attn_mask = attn_mask.float().to(device=self.device)
spec_stacked = spec_stacked.to(device=self.device)
return spec_masked, pos_enc, mask_label, attn_mask, spec_stacked # (x, pos_enc, mask_label, attention_mask. y)
def train(self):
''' Self-Supervised Pre-Training of Transformer Model'''
pbar = tqdm(total=self.total_steps)
while self.global_step <= self.total_steps:
progress = tqdm(self.dataloader, desc="Iteration")
step = 0
loss_val = 0
for batch in progress:
if self.args.online_feat is not None:
# batch are raw waveforms
# batch: (batch_size, channel, max_len)
specs = self.preprocessor(batch.to(device=self.device))
batch = process_train_MAM_data(specs, config=self.transformer_config)
batch_is_valid, *batch = batch
try:
if self.global_step > self.total_steps: break
if not batch_is_valid: continue
step += 1
spec_masked, pos_enc, mask_label, attn_mask, spec_stacked = self.process_data(batch)
loss, pred_spec = self.model(spec_masked, pos_enc, mask_label, attn_mask, spec_stacked)
# Accumulate Loss
if self.gradient_accumulation_steps > 1:
loss = loss / self.gradient_accumulation_steps
if self.apex and self.args.multi_gpu:
raise NotImplementedError
elif self.apex:
self.optimizer.backward(loss)
elif self.args.multi_gpu:
loss = loss.sum()
loss.backward()
else:
loss.backward()
loss_val += loss.item()
# Update
if (step+1) % self.gradient_accumulation_steps == 0:
if self.apex:
# modify learning rate with special warm up BERT uses
# if conifg.apex is False, BertAdam is used and handles this automatically
lr_this_step = self.learning_rate * self.warmup_linear.get_lr(self.global_step, self.warmup_proportion)
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr_this_step
# Step
grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.gradient_clipping)
if math.isnan(grad_norm):
print('[Runner] - Error : grad norm is NaN @ step ' + str(self.global_step))
else:
self.optimizer.step()
self.optimizer.zero_grad()
if self.global_step % self.log_step == 0:
# Log
self.log.add_scalar('lr', self.optimizer.get_lr()[0], self.global_step)
self.log.add_scalar('loss', (loss_val), self.global_step)
self.log.add_scalar('gradient norm', grad_norm, self.global_step)
progress.set_description("Loss %.4f" % (loss_val))
if self.global_step % self.save_step == 0:
self.save_model('states')
mask_spec = self.up_sample_frames(spec_masked[0], return_first=True)
pred_spec = self.up_sample_frames(pred_spec[0], return_first=True)
true_spec = self.up_sample_frames(spec_stacked[0], return_first=True)
mask_spec = plot_spectrogram_to_numpy(mask_spec.data.cpu().numpy())
pred_spec = plot_spectrogram_to_numpy(pred_spec.data.cpu().numpy())
true_spec = plot_spectrogram_to_numpy(true_spec.data.cpu().numpy())
self.log.add_image('mask_spec', mask_spec, self.global_step)
self.log.add_image('pred_spec', pred_spec, self.global_step)
self.log.add_image('true_spec', true_spec, self.global_step)
loss_val = 0
pbar.update(1)
self.global_step += 1
except RuntimeError as e:
if 'CUDA out of memory' in str(e):
print('CUDA out of memory at step: ', self.global_step)
torch.cuda.empty_cache()
self.optimizer.zero_grad()
else:
raise
pbar.close()
self.log.close()