def _test_batched_beam_decoder_step(self, test_args): beam_size = 5 samples, src_dict, tgt_dict = test_utils.prepare_inputs(test_args) num_models = 3 model_list = [] for _ in range(num_models): model_list.append(models.build_model(test_args, src_dict, tgt_dict)) encoder_ensemble = EncoderEnsemble(model_list) # test equivalence # The discrepancy in types here is a temporary expedient. # PyTorch indexing requires int64 while support for tracing # pack_padded_sequence() requires int32. sample = next(samples) src_tokens = sample['net_input']['src_tokens'][0:1].t() src_lengths = sample['net_input']['src_lengths'][0:1].int() pytorch_encoder_outputs = encoder_ensemble(src_tokens, src_lengths) decoder_step_ensemble = DecoderBatchedStepEnsemble( model_list, beam_size=beam_size, ) tmp_dir = tempfile.mkdtemp() decoder_step_pb_path = os.path.join(tmp_dir, 'decoder_step.pb') decoder_step_ensemble.onnx_export( decoder_step_pb_path, pytorch_encoder_outputs, ) # single EOS in flat array input_tokens = torch.LongTensor( np.array([model_list[0].dst_dict.eos()]), ) prev_scores = torch.FloatTensor(np.array([0.0])) timestep = torch.LongTensor(np.array([0])) pytorch_first_step_outputs = decoder_step_ensemble( input_tokens, prev_scores, timestep, *pytorch_encoder_outputs) # next step inputs (input_tokesn shape: [beam_size]) next_input_tokens = torch.LongTensor( np.array([i for i in range(4, 9)]), ) next_prev_scores = pytorch_first_step_outputs[1] next_timestep = timestep + 1 next_states = pytorch_first_step_outputs[4:] step_inputs = [] # encoder outputs need to be replicated for each input hypothesis for encoder_rep in pytorch_encoder_outputs[:len(model_list)]: step_inputs.append(encoder_rep.repeat(1, beam_size, 1)) if model_list[0].decoder.vocab_reduction_module is not None: step_inputs.append(pytorch_encoder_outputs[len(model_list)]) step_inputs.extend(list(next_states)) pytorch_next_step_outputs = decoder_step_ensemble( next_input_tokens, next_prev_scores, next_timestep, *step_inputs) with open(decoder_step_pb_path, 'r+b') as f: onnx_model = onnx.load(f) onnx_decoder = caffe2_backend.prepare(onnx_model) decoder_inputs_numpy = [ next_input_tokens.numpy(), next_prev_scores.detach().numpy(), next_timestep.detach().numpy(), ] for tensor in step_inputs: decoder_inputs_numpy.append(tensor.detach().numpy()) caffe2_next_step_outputs = onnx_decoder.run( tuple(decoder_inputs_numpy), ) for i in range(len(pytorch_next_step_outputs)): caffe2_out_value = caffe2_next_step_outputs[i] pytorch_out_value = pytorch_next_step_outputs[i].data.numpy() np.testing.assert_allclose( caffe2_out_value, pytorch_out_value, rtol=1e-4, atol=1e-6, )
def _test_batched_beam_decoder_step(self, test_args, return_caffe2_rep=False): beam_size = 5 samples, src_dict, tgt_dict = test_utils.prepare_inputs(test_args) task = tasks.DictionaryHolderTask(src_dict, tgt_dict) num_models = 3 model_list = [] for _ in range(num_models): model_list.append(task.build_model(test_args)) encoder_ensemble = EncoderEnsemble(model_list) # test equivalence # The discrepancy in types here is a temporary expedient. # PyTorch indexing requires int64 while support for tracing # pack_padded_sequence() requires int32. sample = next(samples) src_tokens = sample["net_input"]["src_tokens"][0:1].t() src_lengths = sample["net_input"]["src_lengths"][0:1].int() pytorch_encoder_outputs = encoder_ensemble(src_tokens, src_lengths) decoder_step_ensemble = DecoderBatchedStepEnsemble(model_list, tgt_dict, beam_size=beam_size) tmp_dir = tempfile.mkdtemp() decoder_step_pb_path = os.path.join(tmp_dir, "decoder_step.pb") decoder_step_ensemble.onnx_export(decoder_step_pb_path, pytorch_encoder_outputs) # single EOS in flat array input_tokens = torch.LongTensor(np.array([tgt_dict.eos()])) prev_scores = torch.FloatTensor(np.array([0.0])) timestep = torch.LongTensor(np.array([0])) pytorch_first_step_outputs = decoder_step_ensemble( input_tokens, prev_scores, timestep, *pytorch_encoder_outputs) # next step inputs (input_tokesn shape: [beam_size]) next_input_tokens = torch.LongTensor(np.array([i for i in range(4, 9)])) next_prev_scores = pytorch_first_step_outputs[1] next_timestep = timestep + 1 next_states = list(pytorch_first_step_outputs[4:]) # Tile these for the next timestep for i in range(len(model_list)): next_states[i] = next_states[i].repeat(1, beam_size, 1) pytorch_next_step_outputs = decoder_step_ensemble( next_input_tokens, next_prev_scores, next_timestep, *next_states) onnx_decoder = caffe2_backend.prepare_zip_archive(decoder_step_pb_path) if return_caffe2_rep: return onnx_decoder decoder_inputs_numpy = [ next_input_tokens.numpy(), next_prev_scores.detach().numpy(), next_timestep.detach().numpy(), ] for tensor in next_states: decoder_inputs_numpy.append(tensor.detach().numpy()) caffe2_next_step_outputs = onnx_decoder.run( tuple(decoder_inputs_numpy)) for i in range(len(pytorch_next_step_outputs)): caffe2_out_value = caffe2_next_step_outputs[i] pytorch_out_value = pytorch_next_step_outputs[i].detach().numpy() np.testing.assert_allclose(caffe2_out_value, pytorch_out_value, rtol=1e-4, atol=1e-6) decoder_step_ensemble.save_to_db( output_path=os.path.join(tmp_dir, "decoder.predictor_export"), encoder_ensemble_outputs=pytorch_encoder_outputs, )