[egs] support dropout in pytorch training (#3988)

Author: fanlu

egs/aishell/s10/chain/model.py

@@ -13,6 +13,7 @@
 from tdnnf_layer import FactorizedTDNN
 from tdnnf_layer import OrthonormalLinear
 from tdnnf_layer import PrefinalLayer
+from tdnnf_layer import TDNN
 
 
 def get_chain_model(feat_dim,
@@ -101,13 +102,8 @@ def __init__(self,
         num_layers = len(kernel_size_list)
 
         input_dim = feat_dim * 3 + ivector_dim
-        # tdnn1_affine requires [N, T, C]
-        self.tdnn1_affine = nn.Linear(in_features=input_dim,
-                                      out_features=hidden_dim)
-
-        # tdnn1_batchnorm requires [N, C, T]
-        self.tdnn1_batchnorm = nn.BatchNorm1d(num_features=hidden_dim,
-                                              affine=False)
+        
+        self.tdnn1 = TDNN(input_dim=input_dim, hidden_dim=hidden_dim)
 
         tdnnfs = []
         for i in range(num_layers):
@@ -156,7 +152,7 @@ def __init__(self,
 
         self.register_forward_pre_hook(constrain_orthonormal_hook)
 
-    def forward(self, x):
+    def forward(self, x, dropout=0.):
         # input x is of shape: [batch_size, seq_len, feat_dim] = [N, T, C]
         assert x.ndim == 3
 
@@ -178,25 +174,11 @@ def forward(self, x):
 
         # at this point, x is [N, C, T]
 
-        x = x.permute(0, 2, 1)
-
-        # at this point, x is [N, T, C]
-
-        x = self.tdnn1_affine(x)
-
-        # at this point, x is [N, T, C]
-
-        x = F.relu(x)
-
-        x = x.permute(0, 2, 1)
-
-        # at this point, x is [N, C, T]
-
-        x = self.tdnn1_batchnorm(x)
+        x = self.tdnn1(x, dropout=dropout)
 
         # tdnnf requires input of shape [N, C, T]
         for i in range(len(self.tdnnfs)):
-            x = self.tdnnfs[i](x)
+            x = self.tdnnfs[i](x, dropout=dropout)
 
         # at this point, x is [N, C, T]
 

egs/aishell/s10/chain/options.py

@@ -105,6 +105,12 @@ def _set_training_args(parser):
                         dest='l2_regularize',
                         help='l2 regularize',
                         type=float)
+    
+    parser.add_argument('--train.dropout-schedule',
+                        dest='dropout_schedule',
+                        help='dropout schedule',
+                        type=str,
+                        default='0,[email protected],[email protected],0')
 
     parser.add_argument('--train.xent-regularize',
                         dest='xent_regularize',

egs/aishell/s10/chain/tdnnf_layer.py

@@ -51,6 +51,28 @@ def _constrain_orthonormal_internal(M):
     return M
 
 
+class SharedDimScaleDropout(nn.Module):
+    def __init__(self, dim=1):
+        '''
+        Continuous scaled dropout that is const over chosen dim (usually across time)
+        Multiplies inputs by random mask taken from Uniform([1 - 2\alpha, 1 + 2\alpha])
+        '''
+        super().__init__()
+        self.dim = dim
+        self.register_buffer('mask', torch.tensor(0.))
+
+    def forward(self, x, alpha=0.0):
+        if self.training and alpha > 0.:
+            # sample mask from uniform dist with dim of length 1 in self.dim and then repeat to match size
+            tied_mask_shape = list(x.shape)
+            tied_mask_shape[self.dim] = 1
+            repeats = [1 if i != self.dim else x.shape[self.dim]
+                        for i in range(len(x.shape))]
+            return x * self.mask.repeat(tied_mask_shape).uniform_(1 - 2*alpha, 1 + 2*alpha).repeat(repeats)
+            # expected value of dropout mask is 1 so no need to scale outputs like vanilla dropout
+        return x
+
+
 class OrthonormalLinear(nn.Module):
 
     def __init__(self, dim, bottleneck_dim, kernel_size):
@@ -137,6 +159,35 @@ def forward(self, x):
         return x
 
 
+class TDNN(nn.Module):
+    '''
+    This class implements the following topology in kaldi:
+      relu-batchnorm-dropout-layer name=tdnn1 dropout-per-dim-continuous=true dim=1024
+    '''
+
+    def __init__(self, input_dim, hidden_dim):
+        super().__init__()
+        # affine conv1d requires [N, C, T]
+        self.affine = nn.Conv1d(in_channels=input_dim,
+                                out_channels=hidden_dim,
+                                kernel_size=1)
+
+        # tdnn1_batchnorm requires [N, C, T]
+        self.batchnorm = nn.BatchNorm1d(num_features=hidden_dim,
+                                              affine=False)
+        
+        self.dropout = SharedDimScaleDropout(dim=2)
+
+    def forward(self, x, dropout=0.):
+        # input x is of shape: [batch_size, feat_dim, seq_len] = [N, C, T]
+        x = self.affine(x)
+        x = F.relu(x)
+        x = self.batchnorm(x)
+        x = self.dropout(x, alpha=dropout)
+        # return shape is [N, C, T]
+        return x
+
+
 class FactorizedTDNN(nn.Module):
     '''
     This class implements the following topology in kaldi:
@@ -178,7 +229,9 @@ def __init__(self,
         # batchnorm requires [N, C, T]
         self.batchnorm = nn.BatchNorm1d(num_features=dim, affine=False)
 
-    def forward(self, x):
+        self.dropout = SharedDimScaleDropout(dim=2)
+
+    def forward(self, x, dropout=0.):
         # input x is of shape: [batch_size, feat_dim, seq_len] = [N, C, T]
         assert x.ndim == 3
 
@@ -199,9 +252,9 @@ def forward(self, x):
 
         # at this point, x is [N, C, T]
 
-        # TODO(fangjun): implement GeneralDropoutComponent in PyTorch
+        x = self.dropout(x, alpha=dropout)
 
-        if self.linear.kernel_size == 3:
+        if self.linear.kernel_size > 1:
             x = self.bypass_scale * input_x[:, :, self.s:-self.s:self.s] + x
         else:
             x = self.bypass_scale * input_x[:, :, ::self.s] + x

egs/aishell/s10/chain/train.py

@@ -17,6 +17,7 @@
 import torch.optim as optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.nn.utils import clip_grad_value_
+from torch.utils.data.distributed import DistributedSampler
 from torch.utils.tensorboard import SummaryWriter 
 
 import kaldi
@@ -30,10 +31,11 @@
 from common import setup_logger
 from device_utils import allocate_gpu_devices
 from egs_dataset import get_egs_dataloader
+from libs.nnet3.train.dropout_schedule import _get_dropout_proportions
 from model import get_chain_model
 from options import get_args
 
-def get_objf(batch, model, device, criterion, opts, den_graph, training, optimizer=None):
+def get_objf(batch, model, device, criterion, opts, den_graph, training, optimizer=None, dropout=0.):
     total_objf = 0.
     total_weight = 0.
     total_frames = 0.  # for display only
@@ -48,7 +50,7 @@ def get_objf(batch, model, device, criterion, opts, den_graph, training, optimiz
         # at this point, feats is [N, T, C]
         feats = feats.to(device)
         if training:
-            nnet_output, xent_output = model(feats)
+            nnet_output, xent_output = model(feats, dropout=dropout)
         else:
             with torch.no_grad():
                 nnet_output, xent_output = model(feats)
@@ -106,17 +108,20 @@ def get_validation_objf(dataloader, model, device, criterion, opts, den_graph):
     return total_objf, total_weight, total_frames
 
 
-def train_one_epoch(dataloader, valid_dataloader, model, device, optimizer,
-                    criterion, current_epoch, opts, den_graph, tf_writer, rank):
+def train_one_epoch(dataloader, valid_dataloader, model, device, optimizer, criterion, 
+                    current_epoch, num_epochs, opts, den_graph, tf_writer, rank, dropout_schedule):
     total_objf = 0.
     total_weight = 0.
     total_frames = 0.  # for display only
 
     model.train()
-
     for batch_idx, batch in enumerate(dataloader):
+        data_fraction = (batch_idx + current_epoch *
+                         len(dataloader)) / (len(dataloader) * num_epochs)
+        _, dropout = _get_dropout_proportions(
+            dropout_schedule, data_fraction)[0]
         curr_batch_objf, curr_batch_weight, curr_batch_frames = get_objf(
-            batch, model, device, criterion, opts, den_graph, True, optimizer) 
+            batch, model, device, criterion, opts, den_graph, True, optimizer, dropout=dropout)
 
         total_objf += curr_batch_objf
         total_weight += curr_batch_weight
@@ -159,6 +164,11 @@ def train_one_epoch(dataloader, valid_dataloader, model, device, optimizer,
                 'train/current_batch_average_objf',
                 curr_batch_objf / curr_batch_weight,
                 batch_idx + current_epoch * len(dataloader))
+            
+            tf_writer.add_scalar(
+                'train/current_dropout',
+                dropout,
+                batch_idx + current_epoch * len(dataloader))
 
             state_dict = model.state_dict()
             for key, value in state_dict.items():
@@ -205,10 +215,10 @@ def main():
 def process_job(learning_rate, device_id=None, local_rank=None):
     args = get_args()
     if local_rank != None:    
-        setup_logger('{}/train/logs/log-train-rank-{}'.format(args.dir, local_rank),
+        setup_logger('{}/logs/log-train-rank-{}'.format(args.dir, local_rank),
                  args.log_level)
     else:
-        setup_logger('{}/train/logs/log-train-single-GPU'.format(args.dir), args.log_level)
+        setup_logger('{}/logs/log-train-single-GPU'.format(args.dir), args.log_level)
 
     logging.info(' '.join(sys.argv))
 
@@ -249,7 +259,6 @@ def process_job(learning_rate, device_id=None, local_rank=None):
     opts.leaky_hmm_coefficient = args.leaky_hmm_coefficient
 
     den_graph = chain.DenominatorGraph(fst=den_fst, num_pdfs=args.output_dim)
-
 
     model = get_chain_model(
         feat_dim=args.feat_dim,
@@ -325,6 +334,9 @@ def process_job(learning_rate, device_id=None, local_rank=None):
 
             if tf_writer:
                 tf_writer.add_scalar('learning_rate', curr_learning_rate, epoch)
+            
+            if dataloader.sampler and isinstance(dataloader.sampler, DistributedSampler):
+                dataloader.sampler.set_epoch(epoch)
 
             objf = train_one_epoch(dataloader=dataloader,
                                    valid_dataloader=valid_dataloader,
@@ -333,10 +345,12 @@ def process_job(learning_rate, device_id=None, local_rank=None):
                                    optimizer=optimizer,
                                    criterion=criterion,
                                    current_epoch=epoch,
+                                   num_epochs=num_epochs,
                                    opts=opts,
                                    den_graph=den_graph,
                                    tf_writer=tf_writer,
-                                   rank=local_rank)
+                                   rank=local_rank,
+                                   dropout_schedule=args.dropout_schedule)
 
             if best_objf is None:
                 best_objf = objf

egs/aishell/s10/cmd.sh

@@ -13,6 +13,6 @@
 export train_cmd="queue.pl -q all.q --mem 4G"
 export decode_cmd="queue.pl -q all.q --mem 4G"
 export mkgraph_cmd="queue.pl -q all.q --mem 8G"
-export cuda_train_cmd="queue.pl -q v100.q --mem 4G"
-export cuda_inference_cmd="queue.pl -q v100.q --mem 4G"
+export cuda_train_cmd="queue.pl -q g.q --mem 4G"
+export cuda_inference_cmd="queue.pl -q g.q --mem 4G"
 

egs/aishell/s10/local/decode.sh

@@ -11,7 +11,7 @@ lattice_beam=4.0
 max_active=7000 # limit of active tokens
 max_mem=50000000 # approx. limit to memory consumption during minimization in bytes
 min_active=200
-num_threads=20
+num_threads=10
 post_decode_acwt=10  # can be used in 'chain' systems to scale acoustics by 10
 
 . ./path.sh