def generate_attention(model: nn.Module, image_tensor: torch.Tensor, method='transformer_attribution',
                       class_index: int = None, img_size: int = 224, device: str = None) -> np.ndarray:
    """
    Generate attention visualization of the model.
    https://openaccess.thecvf.com/content/CVPR2021/papers/
    Chefer_Transformer_Interpretability_Beyond_Attention_Visualization_CVPR_2021_paper.pdf
    Args:
        model (nn.Module):
            Model to visualize.
        image_tensor (torch.Tensor):
            Image tensor.
        method (str):
            `full`, `rollout` or `transformer_attribution`. Default is `transformer_attribution`.
        class_index (int):
            Class index.
        img_size (int):
            Image size.
        device (str):
            Model device, cpu or cuda

    Returns:
        (`numpy.ndarray`):
            Attention map.
    """
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model.to(device)
    model.eval()
    attribution_generator = LRP(model)
    image_relevance = attribution_generator.generate_lrp(image_tensor.unsqueeze(0).to(device=device),
                                                         method=method,
                                                         index=class_index).detach()
    vis = _reshape_attr_and_get_heatmap(image_relevance, image_tensor, img_size=img_size)
    return vis
Generate attention visualization of the model.
https://openaccess.thecvf.com/content/CVPR2021/papers/
Chefer_Transformer_Interpretability_Beyond_Attention_Visualization_CVPR_2021_paper.pdf
Args:
model (nn.Module):
Model to visualize.
image_tensor (torch.Tensor):
Image tensor.
method (str):
`full`, `rollout` or `transformer_attribution`. Default is `transformer_attribution`.
class_index (int):
Class index.
img_size (int):
Image size.
device (str):
Model device, cpu or cuda

Returns:
(`numpy.ndarray`):
Attention map.
towhee/models/visualization/transformer_visualization.py