| 47 | |
| 48 | |
| 49 | class DecodingWeights(object): |
| 50 | def __init__(self, layer_num, hidden_dim, vocab_size, onmtcheckpoint=None, max_step_for_pe=2048): |
| 51 | self.hidden_dim = hidden_dim |
| 52 | self.max_step_for_pe = max_step_for_pe |
| 53 | # self.w = [] |
| 54 | if onmtcheckpoint: |
| 55 | self.w = {} |
| 56 | for key in onmtcheckpoint: |
| 57 | if key == 'model' or key == 'generator': |
| 58 | self.w[key] = onmtcheckpoint[key] |
| 59 | else: |
| 60 | self.w = {} |
| 61 | self.w['model'] = {} |
| 62 | self.w['generator'] = {} |
| 63 | for i in range(layer_num): |
| 64 | prefix = 'decoder.transformer_layers.' + str(i) |
| 65 | self.w['model'][prefix + '.layer_norm_1.weight'] = torch.zeros(hidden_dim) # self_layernorm_gamma |
| 66 | self.w['model'][prefix + '.layer_norm_1.bias'] = torch.zeros(hidden_dim) # self_layernorm_beta |
| 67 | self.w['model'][prefix + '.self_attn.linear_query.weight'] = torch.zeros(hidden_dim, hidden_dim) # self_kernel_q |
| 68 | self.w['model'][prefix + '.self_attn.linear_keys.weight'] = torch.zeros(hidden_dim, hidden_dim) # self_kernel_k |
| 69 | self.w['model'][prefix + '.self_attn.linear_values.weight'] = torch.zeros(hidden_dim, hidden_dim) # self_kernel_v |
| 70 | self.w['model'][prefix + '.self_attn.linear_query.bias'] = torch.zeros(hidden_dim) # self_bias_q |
| 71 | self.w['model'][prefix + '.self_attn.linear_keys.bias'] = torch.zeros(hidden_dim) # self_bias_k |
| 72 | self.w['model'][prefix + '.self_attn.linear_values.bias'] = torch.zeros(hidden_dim) # self_bias_v |
| 73 | self.w['model'][prefix + '.self_attn.final_linear.weight'] = torch.zeros(hidden_dim, hidden_dim) # self_output_kernel |
| 74 | self.w['model'][prefix + '.self_attn.final_linear.bias'] = torch.zeros(hidden_dim) # self_output_bias |
| 75 | self.w['model'][prefix + '.layer_norm_2.weight'] = torch.zeros(hidden_dim) # cross_layernorm_gamma |
| 76 | self.w['model'][prefix + '.layer_norm_2.bias'] = torch.zeros(hidden_dim) # cross_layernorm_beta |
| 77 | self.w['model'][prefix + '.context_attn.linear_query.weight'] = torch.zeros(hidden_dim, hidden_dim) # cross_kernel_q |
| 78 | self.w['model'][prefix + '.context_attn.linear_keys.weight'] = torch.zeros(hidden_dim, hidden_dim) # cross_kernel_k |
| 79 | self.w['model'][prefix + '.context_attn.linear_values.weight'] = torch.zeros(hidden_dim, hidden_dim) # cross_kernel_v |
| 80 | self.w['model'][prefix + '.context_attn.linear_query.bias'] = torch.zeros(hidden_dim) # cross_bias_q |
| 81 | self.w['model'][prefix + '.context_attn.linear_keys.bias'] = torch.zeros(hidden_dim) # cross_bias_k |
| 82 | self.w['model'][prefix + '.context_attn.linear_values.bias'] = torch.zeros(hidden_dim) # cross_bias_v |
| 83 | self.w['model'][prefix + '.context_attn.final_linear.weight'] = torch.zeros(hidden_dim, hidden_dim) # cross_output_kernel |
| 84 | self.w['model'][prefix + '.context_attn.final_linear.bias'] = torch.zeros(hidden_dim) # cross_output_bias |
| 85 | self.w['model'][prefix + '.feed_forward.layer_norm.weight'] = torch.zeros(hidden_dim) # ffn_layernorm_gamma |
| 86 | self.w['model'][prefix + '.feed_forward.layer_norm.bias'] = torch.zeros(hidden_dim) # ffn_layernorm_beta |
| 87 | self.w['model'][prefix + '.feed_forward.w_1.weight'] = torch.zeros(4 * hidden_dim, hidden_dim) # inter_kernel |
| 88 | self.w['model'][prefix + '.feed_forward.w_1.bias'] = torch.zeros(4 * hidden_dim) # inter_bias |
| 89 | self.w['model'][prefix + '.feed_forward.w_2.weight'] = torch.zeros(hidden_dim, 4 * hidden_dim) # output_kernel |
| 90 | self.w['model'][prefix + '.feed_forward.w_2.bias'] = torch.zeros(hidden_dim) # output_bias |
| 91 | |
| 92 | self.w['model']['decoder.layer_norm.weight'] = torch.zeros(hidden_dim) # decoding_gamma |
| 93 | self.w['model']['decoder.layer_norm.bias'] = torch.zeros(hidden_dim) # decoding_beta |
| 94 | self.w['model']['decoder.embeddings.make_embedding.emb_luts.0.weight'] = torch.zeros(vocab_size, hidden_dim) # embedding_table |
| 95 | |
| 96 | self.w['generator']['0.weight'] = torch.zeros(vocab_size, hidden_dim) |
| 97 | self.w['generator']['0.bias'] = torch.zeros(vocab_size) |
| 98 | |
| 99 | for key in self.w: |
| 100 | if isinstance(self.w[key], dict): |
| 101 | for next_key in self.w[key]: |
| 102 | torch.nn.init.uniform_(self.w[key][next_key], -0.5, 0.5) |
| 103 | else: |
| 104 | torch.nn.init.uniform_(self.w[key], -0.5, 0.5) |
| 105 | |
| 106 |
no outgoing calls
no test coverage detected