| 58 | encoder_example(vars(args)) |
| 59 | |
| 60 | def encoder_example(args): |
| 61 | torch.manual_seed(0) |
| 62 | random.seed(0) |
| 63 | np.random.seed(0) |
| 64 | |
| 65 | batch_size = args['batch_size'] |
| 66 | seq_len = args['seq_len'] |
| 67 | layer_num = args['layer_num'] |
| 68 | head_num = args['head_num'] |
| 69 | head_size = args['head_size'] |
| 70 | hidden_dim = head_num * head_size |
| 71 | |
| 72 | print("\n=============== Argument ===============") |
| 73 | for key in args: |
| 74 | print("{}: {}".format(key, args[key])) |
| 75 | print("========================================\n") |
| 76 | |
| 77 | inp = torch.empty(batch_size, seq_len, hidden_dim).cuda() |
| 78 | torch.nn.init.normal_(inp, -0.02, 0.02) |
| 79 | mem_seq_lens = torch.randint(1, seq_len+1, (batch_size,), dtype=torch.int32).cuda() |
| 80 | if args['remove_padding']: |
| 81 | if args['avg_seq_len'] > 0: |
| 82 | mem_seq_lens = torch.ones((batch_size,)) * args['avg_seq_len'] |
| 83 | mem_seq_lens = mem_seq_lens.to(torch.int32).cuda() |
| 84 | elif args['avg_seq_len'] == -1: |
| 85 | mem_seq_lens = torch.ones((batch_size,)) * seq_len / 2 |
| 86 | mem_seq_lens = mem_seq_lens.to(torch.int32).cuda() |
| 87 | else: |
| 88 | raise ValueError("wrong avg_seq_len") |
| 89 | |
| 90 | mask = ~sequence_mask(mem_seq_lens, seq_len).unsqueeze(1) |
| 91 | if args['data_type'] == 'fp16': |
| 92 | inp = inp.half() |
| 93 | |
| 94 | weights = EncoderWeights(layer_num, hidden_dim) |
| 95 | |
| 96 | onmt_encoder = ONMTEncoder(layer_num, hidden_dim, head_num, 4 * hidden_dim, weights) |
| 97 | onmt_encoder.cuda() |
| 98 | if args['data_type'] == 'fp16': |
| 99 | onmt_encoder.half() |
| 100 | onmt_encoder.eval() |
| 101 | onmt_encoder = torch.jit.trace(onmt_encoder, (inp, mask)) |
| 102 | |
| 103 | if args['data_type'] == 'fp16': |
| 104 | weights.to_half() |
| 105 | weights.to_cuda() |
| 106 | custom_encoder = CustomEncoder(layer_num, head_num, head_size, weights, |
| 107 | remove_padding=False, allow_gemm_test=args['allow_gemm_test'], |
| 108 | path=args['ths_path']) |
| 109 | custom_encoder = torch.jit.script(custom_encoder) |
| 110 | |
| 111 | eff_custom_encoder = CustomEncoder(layer_num, head_num, head_size, weights, |
| 112 | remove_padding=True, allow_gemm_test=args['allow_gemm_test'], |
| 113 | path=args['ths_path']) |
| 114 | eff_custom_encoder = torch.jit.script(eff_custom_encoder) |
| 115 | |
| 116 | with torch.no_grad(): |
| 117 | output_mask = sequence_mask(mem_seq_lens, args['seq_len']).to(mask.dtype).unsqueeze(-1) |