TransformerでのSelf-Attentionの可視化

現在、fairseqというキットを用いて機械翻訳のモデルをいじっています。
そのモデルに対して、学習時と推論時に各テキストごとにどの様なself-attentionが張られているのか知りたいです。

複数のクラスのプログラムの中で、どの部分でattentionを出力すればself-attentionが求められるのか教えていただけるでしょうか?

最終的にはself-attentinoのスコアを元に以下のようなヒートマップを作りたいと考えています。