AF3 AtomAttentionDecoder类源码解读

AlphaFold3的AtomAttentionDecoder类旨在从每个 token 的表示扩展到每个原子的表示，同时通过交错注意力机制对原子及其对关系举行建模。这种设计可以在生物分子建模中捕获复杂的原子级别交互。
源代码：

1. class AtomAttentionDecoder(nn.Module):
2.     """AtomAttentionDecoder that broadcasts per-token activations to per-atom activations."""
4.     def __init__(
5.             self,
6.             c_token: int,
7.             c_atom: int = 128,
8.             c_atompair: int = 16,
9.             no_blocks: int = 3,
10.             no_heads: int = 8,
11.             dropout=0.0,
12.             n_queries: int = 32,
13.             n_keys: int = 128,
14.             clear_cache_between_blocks: bool = False
15.     ):
16.         """Initialize the AtomAttentionDecoder module.
17.         Args:
18.             c_token:
19.                 The number of channels for the token representation.
20.             c_atom:
21.                 The number of channels for the atom representation. Defaults to 128.
22.             c_atompair:
23.                 The number of channels for the atom pair representation. Defaults to 16.
24.             no_blocks:
25.                 Number of blocks.
26.             no_heads:
27.                 Number of parallel attention heads. Note that c_atom will be split across no_heads
28.                 (i.e. each head will have dimension c_atom // no_heads).
29.             dropout:
30.                 Dropout probability on attn_output_weights. Default: 0.0 (no dropout).
31.             n_queries:
32.                 The size of the atom window. Defaults to 32.
33.             n_keys:
34.                 Number of atoms each atom attends to in local sequence space. Defaults to 128.
35.             clear_cache_between_blocks:
36.                 Whether to clear CUDA's GPU memory cache between blocks of the
37.                 stack. Slows down each block but can reduce fragmentation
38.         """
39.         super().__init__()
40.         self.c_token = c_token
41.         self.c_atom = c_atom
42.         self.c_atompair = c_atompair
43.         self.num_blocks = no_blocks
44.         self.num_heads = no_heads
45.         self.dropout = dropout
46.         self.n_queries = n_queries
47.         self.n_keys = n_keys
48.         self.clear_cache_bet

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。