(self)
Setup ASR model and optimizer
Setup ASR model and optimizer
Setup ASR model and optimizer
def set_model(self):
        ''' Setup ASR model and optimizer '''

        # Model
        self.model = RNNLM(self.vocab_size, **
                           self.config['model']).to(self.device)
        self.verbose(self.model.create_msg())
        # Losses
        self.seq_loss = torch.nn.CrossEntropyLoss(ignore_index=0)
        # Optimizer
        self.optimizer = Optimizer(
            self.model.parameters(), **self.config['hparas'])
        # Enable AMP if needed
        self.enable_apex()
        # load pre-trained model
        if self.paras.load:
            self.load_ckpt()
            ckpt = torch.load(self.paras.load, map_location=self.device)
            self.model.load_state_dict(ckpt['model'])
            self.optimizer.load_opt_state_dict(ckpt['optimizer'])
            self.step = ckpt['global_step']
            self.verbose('Load ckpt from {}, restarting at step {}'.format(
                self.paras.load, self.step))
def set_model ( self ) : # Model self . model = RNNLM ( self . vocab_size , * * self . config [ 'model' ] ) . to ( self . device ) self . verbose ( self . model . create_msg ( ) ) # Losses self . seq_loss = torch . nn . CrossEntropyLoss ( ignore_index = 0 ) # Optimizer self . optimizer = Optimizer ( self . model . parameters ( ) , * * self . config [ 'hparas' ] ) # Enable AMP if needed self . enable_apex ( ) # load pre-trained model if self . paras . load : self . load_ckpt ( ) ckpt = torch . load ( self . paras . load , map_location = self . device ) self . model . load_state_dict ( ckpt [ 'model' ] ) self . optimizer . load_opt_state_dict ( ckpt [ 'optimizer' ] ) self . step = ckpt [ 'global_step' ] self . verbose ( 'Load ckpt from {}, restarting at step {}' . format ( self . paras . load , self . step ) )