| 55 | return cache |
| 56 | |
| 57 | class ONMTDecoder(torch.nn.Module): |
| 58 | def __init__(self, layer_num, head_num, head_size, weights): |
| 59 | super().__init__() |
| 60 | self.layer_num = layer_num |
| 61 | self.hidden_dim = head_num * head_size |
| 62 | self.decoders = torch.nn.ModuleList() |
| 63 | for i in range(layer_num): |
| 64 | self.decoders.append(TransformerDecoderLayer(self.hidden_dim, head_num, 4 * self.hidden_dim, 0, 0)) |
| 65 | for i in range(layer_num): |
| 66 | prefix = 'decoder.transformer_layers.' + str(i) |
| 67 | self.decoders[i].layer_norm_1.weight.data = weights.w['model'][prefix + '.layer_norm_1.weight'] |
| 68 | self.decoders[i].layer_norm_1.bias.data = weights.w['model'][prefix + '.layer_norm_1.bias'] |
| 69 | self.decoders[i].self_attn.linear_query.weight.data = weights.w['model'][prefix + '.self_attn.linear_query.weight'] |
| 70 | self.decoders[i].self_attn.linear_keys.weight.data = weights.w['model'][prefix + '.self_attn.linear_keys.weight'] |
| 71 | self.decoders[i].self_attn.linear_values.weight.data = weights.w['model'][prefix + '.self_attn.linear_values.weight'] |
| 72 | self.decoders[i].self_attn.linear_query.bias.data = weights.w['model'][prefix + '.self_attn.linear_query.bias'] |
| 73 | self.decoders[i].self_attn.linear_keys.bias.data = weights.w['model'][prefix + '.self_attn.linear_keys.bias'] |
| 74 | self.decoders[i].self_attn.linear_values.bias.data = weights.w['model'][prefix + '.self_attn.linear_values.bias'] |
| 75 | self.decoders[i].self_attn.final_linear.weight.data = weights.w['model'][prefix + '.self_attn.final_linear.weight'] |
| 76 | self.decoders[i].self_attn.final_linear.bias.data = weights.w['model'][prefix + '.self_attn.final_linear.bias'] |
| 77 | self.decoders[i].layer_norm_2.weight.data = weights.w['model'][prefix + '.layer_norm_2.weight'] |
| 78 | self.decoders[i].layer_norm_2.bias.data = weights.w['model'][prefix + '.layer_norm_2.bias'] |
| 79 | self.decoders[i].context_attn.linear_query.weight.data = weights.w['model'][prefix + '.context_attn.linear_query.weight'] |
| 80 | self.decoders[i].context_attn.linear_keys.weight.data = weights.w['model'][prefix + '.context_attn.linear_keys.weight'] |
| 81 | self.decoders[i].context_attn.linear_values.weight.data = weights.w['model'][prefix + '.context_attn.linear_values.weight'] |
| 82 | self.decoders[i].context_attn.linear_query.bias.data = weights.w['model'][prefix + '.context_attn.linear_query.bias'] |
| 83 | self.decoders[i].context_attn.linear_keys.bias.data = weights.w['model'][prefix + '.context_attn.linear_keys.bias'] |
| 84 | self.decoders[i].context_attn.linear_values.bias.data = weights.w['model'][prefix + '.context_attn.linear_values.bias'] |
| 85 | self.decoders[i].context_attn.final_linear.weight.data = weights.w['model'][prefix + '.context_attn.final_linear.weight'] |
| 86 | self.decoders[i].context_attn.final_linear.bias.data = weights.w['model'][prefix + '.context_attn.final_linear.bias'] |
| 87 | self.decoders[i].feed_forward.layer_norm.weight.data = weights.w['model'][prefix + '.feed_forward.layer_norm.weight'] |
| 88 | self.decoders[i].feed_forward.layer_norm.bias.data = weights.w['model'][prefix + '.feed_forward.layer_norm.bias'] |
| 89 | self.decoders[i].feed_forward.w_1.weight.data = weights.w['model'][prefix + '.feed_forward.w_1.weight'] |
| 90 | self.decoders[i].feed_forward.w_1.bias.data = weights.w['model'][prefix + '.feed_forward.w_1.bias'] |
| 91 | self.decoders[i].feed_forward.w_2.weight.data = weights.w['model'][prefix + '.feed_forward.w_2.weight'] |
| 92 | self.decoders[i].feed_forward.w_2.bias.data = weights.w['model'][prefix + '.feed_forward.w_2.bias'] |
| 93 | |
| 94 | def forward(self, inputs, memory, src_pad_msk, cache, step): |
| 95 | output = inputs |
| 96 | for i in range(self.layer_num): |
| 97 | output, _, _ = self.decoders[i](output, memory, src_pad_msk, None, cache[i], step) |
| 98 | return output |