Method _get_visual_features

axlearn/vision/virtex.py:116–155

Gets the visual features. The visual features will be extracted from a NestedTensor based on a path specified in `cfg.visual_feature_layer_name`. The NestedTensor can be either `visual_outputs` or self.get_module_outputs()["visual"]. Args: visual_output

(self, visual_outputs: NestedTensor)

        return self.embed_image(input_batch["image"])

    def _get_visual_features(self, visual_outputs: NestedTensor) -> Tensor:
        """Gets the visual features.

        The visual features will be extracted from a NestedTensor based on a path specified in
        `cfg.visual_feature_layer_name`.

        The NestedTensor can be either `visual_outputs` or self.get_module_outputs()["visual"].

        Args:
            visual_outputs: The outputs of the visual layer.

        Returns:
            visual_features: The features from the visual backbone.
                Shape: (batch, height, width, channels)

        Raises:
            ValueError: If visual_feature_layer_name cannot be found.
        """
        cfg = self.config
        try:
            return get_recursively(visual_outputs, cfg.visual_feature_layer_name)
        except KeyError:
            pass

        try:
            return get_recursively(
                self.get_module_outputs(),
                f"visual/{cfg.visual_feature_layer_name}",
            )
        except KeyError:
            pass

        def _paths(x):
            return jax.tree_util.tree_leaves(tree_paths(x))

        raise ValueError(
            f"Cannot find visual features at {cfg.visual_feature_layer_name}. "
            f"visual_outputs={_paths(visual_outputs)}, "
            f"module_outputs={_paths(self.get_module_outputs().get('visual'))}"
        )


class VirTexModel(ImageBackboneModelMixin, BaseLayer):