def parallelize()

in optimum/graphcore/models/lxmert/modeling_lxmert.py [0:0]

• 31 lines of code
• 7 McCabe index (conditional complexity)

    def parallelize(self):
        """
        Transform the model to run in an IPU pipeline.
        - Adds pipeline stages to the model
        - Adds recomputation checkpoints
        Recommended usage:
        ```
        model = PipelinedLxmertForQuestionAnswering(config).parallelize().half()
        ```
        """
        self._hooks = []

        logger.info("-------------------- Device Allocation --------------------")
        logger.info("Embedding               --> IPU 0")
        self.lxmert.embeddings = poptorch.BeginBlock(self.lxmert.embeddings, "Embedding", ipu_id=0)
        logger.info("Image embedding         --> IPU 0")
        self.lxmert.encoder.visn_fc = poptorch.BeginBlock(self.lxmert.encoder.visn_fc, "Image embedding", ipu_id=0)

        # Language layers
        for index, layer in enumerate(self.lxmert.encoder.layer):
            if self.ipu_config.recompute_checkpoint_every_layer:
                h = recomputation_checkpoint(layer)
                self._hooks.append(h)
            self.lxmert.encoder.layer[index] = poptorch.BeginBlock(layer, f"Language layer{index}", ipu_id=1)
            logger.info(f"Language layer {index:<2}       --> IPU 1")

        # Visual layers
        for index, layer in enumerate(self.lxmert.encoder.r_layers):
            if self.ipu_config.recompute_checkpoint_every_layer:
                h = recomputation_checkpoint(layer)
                self._hooks.append(h)
            self.lxmert.encoder.r_layers[index] = poptorch.BeginBlock(layer, f"Visual layer{index}", ipu_id=2)
            logger.info(f"Visual layer {index:<2}         --> IPU 2")

        # Cross modality layers
        for index, layer in enumerate(self.lxmert.encoder.x_layers):
            if self.ipu_config.recompute_checkpoint_every_layer:
                h = recomputation_checkpoint(layer)
                self._hooks.append(h)
            self.lxmert.encoder.x_layers[index] = poptorch.BeginBlock(layer, f"Cross modality layer{index}", ipu_id=3)
            logger.info(f"Cross modality layer {index:<2} --> IPU 3")

        logger.info("Pooler                  --> IPU 3")
        self.lxmert.pooler = poptorch.BeginBlock(self.lxmert.pooler, "Pooler", ipu_id=3)

        logger.info("Head                    --> IPU 3")
        self.answer_head = poptorch.BeginBlock(self.answer_head, "Head", ipu_id=3)
        logger.info("-----------------------------------------------------------")
        return self