| 90 | |
| 91 | |
| 92 | class CustomEncoder(torch.nn.Module): |
| 93 | def __init__(self, layer_num, head_num, head_size, weights, |
| 94 | remove_padding=False, allow_gemm_test=False, path='./lib/libth_transformer.so', embedding=None): |
| 95 | super().__init__() |
| 96 | self.layer_num = layer_num |
| 97 | self.remove_padding = remove_padding |
| 98 | self.embedding = embedding |
| 99 | torch.classes.load_library(path) |
| 100 | |
| 101 | weights_ = weights.listed_weights() |
| 102 | assert len(weights_) == 18 |
| 103 | try: |
| 104 | self.encoders = torch.classes.FasterTransformer.Encoder( |
| 105 | *weights_, |
| 106 | head_num, head_size, 4 * head_num * head_size, remove_padding, layer_num, allow_gemm_test, False, 1.0) |
| 107 | except: |
| 108 | # legacy ths for 20.03 image |
| 109 | self.encoders = torch.classes.FasterTransformerEncoder( |
| 110 | *weights_, |
| 111 | head_num, head_size, 4 * head_num * head_size, remove_padding, layer_num, allow_gemm_test, False, 1.0) |
| 112 | |
| 113 | def forward(self, inputs, lengths): |
| 114 | if self.embedding is not None: |
| 115 | emb = self.embedding(inputs) |
| 116 | inputs = emb.transpose(0, 1).contiguous() |
| 117 | hidden_states = self.encoders.forward(inputs, lengths) |
| 118 | if self.embedding is not None: |
| 119 | return emb, hidden_states.transpose(0, 1).contiguous(), lengths |
| 120 | else: |
| 121 | return hidden_states |
| 122 | |
| 123 | |
| 124 | class ONMTEncoder(torch.nn.Module): |
no outgoing calls
no test coverage detected