Class CustomEncoder

examples/pytorch/bert/utils/encoder.py:270–324

class CustomEncoder(torch.nn.Module):
    def __init__(self, layer_num, head_num, head_size, weights,
                 int8_mode=0, remove_padding=False, sparse=False,
                 path='./lib/libth_transformer.so', tensor_para_size=1,
                 pipeline_para_size=1):
        super().__init__()
        self.layer_num = layer_num
        self.remove_padding = remove_padding
        self.int8_mode = int8_mode
        torch.classes.load_library(path)

        weights_ = weights.listed_weights()

        self.use_mpi = dist.is_mpi_available()

        if self.use_mpi:
            try:
                dist.init_process_group(backend='mpi')
            except:
                print("[INFO] WARNING: Exception occurred in dist.init_process_group(backend='mpi'). Maybe the process group has been initialized somewhere else.")
        else:
            print("[INFO] MPI is not available in this PyTorch build.")
            assert tensor_para_size == 1, "[FATAL] MPI is required for tensor_para_size > 1."
            assert pipeline_para_size == 1, "[FATAL] MPI is required for pipeline_para_size > 1."

        if int8_mode == 0:
            assert len(weights_) == 16
            try:
                self.encoders = torch.classes.FasterTransformer.Bert(
                    *weights_,
                    head_num, head_size, 4 * head_num * head_size, remove_padding, layer_num, sparse, 1.0,
                    tensor_para_size, pipeline_para_size)
            except:
                # legacy ths for 20.03 image
                self.encoders = torch.classes.FasterTransformerBert(
                    *weights_,
                    head_num, head_size, 4 * head_num * head_size, remove_padding, layer_num, sparse, 1.0,
                    tensor_para_size, pipeline_para_size)
        else:
            assert len(weights_) == 18
            assert tensor_para_size == 1, "INT8 BERT still only support tensor_para_size = 1"
            assert pipeline_para_size == 1, "INT8 BERT still only support pipeline_para_size = 1"
            try:
                self.encoders = torch.classes.FasterTransformer.INT8Bert(
                    *weights_,
                    head_num, head_size, remove_padding, layer_num, int8_mode, sparse, 1.0)
            except:
                # legacy ths for 20.03 image
                self.encoders = torch.classes.FasterTransformerINT8Bert(
                    *weights_,
                    head_num, head_size, remove_padding, layer_num, int8_mode, sparse, 1.0)

    def forward(self, hidden_states, attention_mask, sequence_lengths):
        hidden_states = self.encoders.forward(hidden_states, sequence_lengths)
        return (hidden_states,)


class HuggingFaceEncoder(torch.nn.Module):