Merge pull request #251 from Xreki/opt/einsum_slice

Optimize the use of getitem (slice) and einsum.

Author: SuperXiang

apps/protein_folding/helixfold/alphafold_paddle/model/all_atom.py

@@ -64,7 +64,8 @@ def get_chi_atom_indices():
             atom_indices.append(
                 [residue_constants.atom_order[atom] for atom in chi_angle])
         for _ in range(4 - len(atom_indices)):
-            atom_indices.append([0, 0, 0, 0])  # For chi angles not defined on the AA.
+            atom_indices.append(
+                [0, 0, 0, 0])  # For chi angles not defined on the AA.
         chi_atom_indices.append(atom_indices)
 
     chi_atom_indices.append([[0, 0, 0, 0]] * 4)  # For UNKNOWN residue.
@@ -274,8 +275,7 @@ def atom37_to_torsion_angles(
     aatype: paddle.Tensor,  # (B, T, N)
     all_atom_pos: paddle.Tensor,  # (B, T, N, 37, 3)
     all_atom_mask: paddle.Tensor,  # (B, T, N, 37)
-    placeholder_for_undefined=False,
-) -> Dict[str, paddle.Tensor]:
+    placeholder_for_undefined=False, ) -> Dict[str, paddle.Tensor]:
     """Computes the 7 torsion angles (in sin, cos encoding) for each residue.
 
     The 7 torsion angles are in the order
@@ -300,44 +300,57 @@ def atom37_to_torsion_angles(
     """
 
     # Map aatype > 20 to 'Unknown' (20).
-    aatype = paddle.minimum(aatype.astype('int'), paddle.to_tensor([20]).astype('int'))
-    
+    aatype = paddle.minimum(
+        aatype.astype('int'), paddle.full(shape=[1], fill_value=20, dtype="int"))
+
     num_batch, num_temp, num_res = aatype.shape
 
     # Compute the backbone angles.
     pad = paddle.zeros([num_batch, num_temp, 1, 37, 3])
-    prev_all_atom_pos = paddle.concat([pad, all_atom_pos[..., :-1, :, :]], axis=-3)
+    prev_all_atom_pos = paddle.concat(
+        [pad, all_atom_pos[..., :-1, :, :]], axis=-3)
 
     pad = paddle.zeros([num_batch, num_temp, 1, 37])
-    prev_all_atom_mask = paddle.concat([pad, all_atom_mask[..., :-1, :]], axis=-2)
+    prev_all_atom_mask = paddle.concat(
+        [pad, all_atom_mask[..., :-1, :]], axis=-2)
 
     # For each torsion angle collect the 4 atom positions that define this angle.
     # shape (B, T, N, atoms=4, xyz=3)
     pre_omega_atom_pos = paddle.concat(
-        [prev_all_atom_pos[..., 1:3, :],  # prev CA, C
-        all_atom_pos[..., 0:2, :]  # this N, CA
-        ], axis=-2)
+        [
+            prev_all_atom_pos[..., 1:3, :],  # prev CA, C
+            all_atom_pos[..., 0:2, :]  # this N, CA
+        ],
+        axis=-2)
 
     phi_atom_pos = paddle.concat(
-        [prev_all_atom_pos[..., 2:3, :],  # prev C
-        all_atom_pos[..., 0:3, :]  # this N, CA, C
-        ], axis=-2)
+        [
+            prev_all_atom_pos[..., 2:3, :],  # prev C
+            all_atom_pos[..., 0:3, :]  # this N, CA, C
+        ],
+        axis=-2)
 
     psi_atom_pos = paddle.concat(
-        [all_atom_pos[..., 0:3, :],  # this N, CA, C
-        all_atom_pos[..., 4:5, :]  # this O
-        ], axis=-2)
+        [
+            all_atom_pos[..., 0:3, :],  # this N, CA, C
+            all_atom_pos[..., 4:5, :]  # this O
+        ],
+        axis=-2)
 
     # Collect the masks from these atoms.
     # Shape [batch, n_temp, num_res]
     pre_omega_mask = (
-        paddle.prod(prev_all_atom_mask[..., 1:3], axis=-1)  # prev CA, C
-        * paddle.prod(all_atom_mask[..., 0:2], axis=-1))  # this N, CA
+        paddle.prod(
+            prev_all_atom_mask[..., 1:3], axis=-1)  # prev CA, C
+        * paddle.prod(
+            all_atom_mask[..., 0:2], axis=-1))  # this N, CA
     phi_mask = (
         prev_all_atom_mask[..., 2]  # prev C
-        * paddle.prod(all_atom_mask[..., 0:3], axis=-1))  # this N, CA, C
+        * paddle.prod(
+            all_atom_mask[..., 0:3], axis=-1))  # this N, CA, C
     psi_mask = (
-        paddle.prod(all_atom_mask[..., 0:3], axis=-1) *  # this N, CA, C
+        paddle.prod(
+            all_atom_mask[..., 0:3], axis=-1) *  # this N, CA, C
         all_atom_mask[..., 4])  # this O
 
     # Collect the atoms for the chi-angles.
@@ -375,18 +388,18 @@ def atom37_to_torsion_angles(
     # Stack all torsion angle atom positions.
     # Shape (B, T, N, torsions=7, atoms=4, xyz=3)
     torsions_atom_pos = paddle.concat(
-        [pre_omega_atom_pos[:, :, :, None, :, :],
-        phi_atom_pos[:, :, :, None, :, :],
-        psi_atom_pos[:, :, :, None, :, :],
+        [pre_omega_atom_pos.unsqueeze(axis=-3), # [:, :, :, None, :, :]
+        phi_atom_pos.unsqueeze(axis=-3), # [:, :, :, None, :, :]
+        psi_atom_pos.unsqueeze(axis=-3), # [:, :, :, None, :, :]
         chis_atom_pos
         ], axis=3)
 
     # Stack up masks for all torsion angles.
     # shape (B, T, N, torsions=7)
     torsion_angles_mask = paddle.concat(
-        [pre_omega_mask[..., None],
-        phi_mask[..., None],
-        psi_mask[..., None],
+        [pre_omega_mask.unsqueeze(axis=-1), # [..., None]
+        phi_mask.unsqueeze(axis=-1), # [..., None]
+        psi_mask.unsqueeze(axis=-1), # [..., None]
         chis_mask
         ], axis=-1)
 
@@ -417,7 +430,7 @@ def atom37_to_torsion_angles(
 
     # Mirror psi, because we computed it from the Oxygen-atom.
     torsion_angles_sin_cos *= paddle.to_tensor(
-        [1., 1., -1., 1., 1., 1., 1.])[None, None, None, :, None]
+        [1., 1., -1., 1., 1., 1., 1.]).reshape([1, 1, 1, 7, 1]) # [None, None, None, :, None]
 
     # Create alternative angles for ambiguous atom names.
     chi_is_ambiguous = utils.batched_gather(
@@ -428,7 +441,7 @@ def atom37_to_torsion_angles(
         1.0 - 2.0 * chi_is_ambiguous], axis=-1)
     # mirror_torsion_angles (B, T, N, torsions=7)
     alt_torsion_angles_sin_cos = (
-        torsion_angles_sin_cos * mirror_torsion_angles[:, :, :, :, None])
+        torsion_angles_sin_cos * mirror_torsion_angles.unsqueeze(axis=-1))
 
     if placeholder_for_undefined:
         # Add placeholder torsions in place of undefined torsion angles
@@ -437,10 +450,8 @@ def atom37_to_torsion_angles(
             paddle.ones(torsion_angles_sin_cos.shape[:-1]),
             paddle.zeros(torsion_angles_sin_cos.shape[:-1])
         ], axis=-1)
-        torsion_angles_sin_cos = torsion_angles_sin_cos * torsion_angles_mask[
-            ..., None] + placeholder_torsions * (1 - torsion_angles_mask[..., None])
-        alt_torsion_angles_sin_cos = alt_torsion_angles_sin_cos * torsion_angles_mask[
-            ..., None] + placeholder_torsions * (1 - torsion_angles_mask[..., None])
+        torsion_angles_sin_cos = torsion_angles_sin_cos * torsion_angles_mask.unsqueeze(axis=-1) + placeholder_torsions * (1 - torsion_angles_mask.unsqueeze(axis=-1))
+        alt_torsion_angles_sin_cos = alt_torsion_angles_sin_cos * torsion_angles_mask.unsqueeze(axis=-1) + placeholder_torsions * (1 - torsion_angles_mask.unsqueeze(axis=-1))
 
     return {
         'torsion_angles_sin_cos': torsion_angles_sin_cos,  # (B, T, N, 7, 2)
@@ -579,7 +590,7 @@ def frames_and_literature_positions_to_atom14_pos(
     """
     # Pick the appropriate transform for every atom.
     restype_atom14_to_rigid_group = paddle.to_tensor(
-        residue_constants.restype_atom14_to_rigid_group)[None, ...]
+        residue_constants.restype_atom14_to_rigid_group).unsqueeze(axis=0)
 
     # [1, 21, 14] -> # [n_batch, 21, 14]
     n_batch = aatype.shape[0]
@@ -612,7 +623,7 @@ def _convert(x, y):
     # Gather the literature atom positions for each residue.
     # r3.Vecs with shape (B, N, 14)
     restype_atom14_rigid_group_positions = paddle.to_tensor(
-        residue_constants.restype_atom14_rigid_group_positions)[None, ...]
+        residue_constants.restype_atom14_rigid_group_positions).unsqueeze(axis=0)
     # [1, 21, 14, 3] -> [B, 21, 14, 3]
     if n_batch > 1:
         restype_atom14_rigid_group_positions = paddle.tile(
@@ -629,7 +640,7 @@ def _convert(x, y):
 
     # Mask out non-existing atoms.
     restype_atom14_mask = paddle.to_tensor(
-        residue_constants.restype_atom14_mask)[None, ...]
+        residue_constants.restype_atom14_mask).unsqueeze(axis=0)
     # [1, 21, 14] -> [B, 21, 14]
     if n_batch > 1:
         restype_atom14_mask = paddle.tile(

apps/protein_folding/helixfold/alphafold_paddle/model/config.py

@@ -404,7 +404,9 @@ def model_config(name: str) -> ml_collections.ConfigDict:
             'use_remat': False,
             'zero_init': True,
             'low_memory': False,
+            'fuse_linear': False,
             'fuse_attention': True,
+            'use_flash_attn': True,
             'use_dropout_nd': True,
             'outer_product_mean_position': 'origin', # 'origin' or 'middle', 'first', 'end', set 'end' if use BP
         },

apps/protein_folding/helixfold/alphafold_paddle/model/folding.py

@@ -51,6 +51,8 @@ def __init__(self, channel_num, config, global_config,
         self.global_config = global_config
         self.dist_epsilon = dist_epsilon
 
+        Linear = paddle.incubate.nn.FusedLinear if self.global_config.fuse_linear else paddle.nn.Linear
+
         num_head = self.config.num_head
         num_scalar_qk = self.config.num_scalar_qk
         num_point_qk = self.config.num_point_qk
@@ -62,15 +64,15 @@ def __init__(self, channel_num, config, global_config,
         assert num_point_qk > 0
         assert num_point_v > 0
 
-        self.q_scalar = nn.Linear(
+        self.q_scalar = Linear(
             channel_num['seq_channel'], num_head * num_scalar_qk)
-        self.kv_scalar = nn.Linear(
+        self.kv_scalar = Linear(
             channel_num['seq_channel'],
             num_head * (num_scalar_v + num_scalar_qk))
 
-        self.q_point_local = nn.Linear(
+        self.q_point_local = Linear(
             channel_num['seq_channel'], num_head * 3 * num_point_qk)
-        self.kv_point_local = nn.Linear(
+        self.kv_point_local = Linear(
             channel_num['seq_channel'],
             num_head * 3 * (num_point_qk + num_point_v))
 
@@ -79,15 +81,15 @@ def __init__(self, channel_num, config, global_config,
             [num_head], 'float32',
             default_initializer=nn.initializer.Constant(tpw))
 
-        self.attention_2d = nn.Linear(channel_num['pair_channel'], num_head)
+        self.attention_2d = Linear(channel_num['pair_channel'], num_head)
 
         if self.global_config.zero_init:
             init_w = nn.initializer.Constant(value=0.0)
         else:
             init_w = nn.initializer.XavierUniform()
 
         c = num_scalar_v + num_point_v * 4 + channel_num['pair_channel']
-        self.output_projection = nn.Linear(
+        self.output_projection = Linear(
             num_head * c, num_output,
             weight_attr=paddle.ParamAttr(initializer=init_w))
 
@@ -257,6 +259,8 @@ def __init__(self, channel_num, config, global_config):
         self.config = config
         self.global_config = global_config
 
+        Linear = paddle.incubate.nn.FusedLinear if self.global_config.fuse_linear else paddle.nn.Linear
+
         self.invariant_point_attention = InvariantPointAttention(
             channel_num, config, global_config)
         self.attention_layer_norm = nn.LayerNorm(channel_num['seq_channel'])
@@ -273,7 +277,7 @@ def __init__(self, channel_num, config, global_config):
             if i > 0:
                 layer_name, c_in = f'transition_{i}', self.config.num_channel
 
-            setattr(self, layer_name, nn.Linear(
+            setattr(self, layer_name, Linear(
                 c_in, self.config.num_channel,
                 weight_attr=paddle.ParamAttr(initializer=init_w)))
 
@@ -287,7 +291,7 @@ def __init__(self, channel_num, config, global_config):
             last_init_w = nn.initializer.XavierUniform()
 
         # Jumper et al. (2021) Alg. 23 "Backbone update"
-        self.affine_update = nn.Linear(
+        self.affine_update = Linear(
             self.config.num_channel, 6,
             weight_attr=paddle.ParamAttr(initializer=last_init_w))
 
@@ -349,8 +353,10 @@ def __init__(self, channel_num, config, global_config):
         self.config = config
         self.global_config = global_config
 
+        Linear = paddle.incubate.nn.FusedLinear if self.global_config.fuse_linear else paddle.nn.Linear
+
         self.single_layer_norm = nn.LayerNorm(channel_num['seq_channel'])
-        self.initial_projection = nn.Linear(
+        self.initial_projection = Linear(
             channel_num['seq_channel'], config.num_channel)
         self.pair_layer_norm = nn.LayerNorm(channel_num['pair_channel'])
 
@@ -861,10 +867,10 @@ def supervised_chi_loss(ret, batch, value, config):
 
     residue_type_one_hot = paddle.nn.functional.one_hot(batch['aatype_index'], 
                             num_classes=residue_constants.restype_num + 1)
-    chi_pi_periodic = paddle.einsum('nijk,nkl->nijl', residue_type_one_hot[:, None, ...], 
+    chi_pi_periodic = paddle.einsum('nijk,nkl->nijl', residue_type_one_hot.unsqueeze(axis=1),
                             paddle.to_tensor(residue_constants.chi_pi_periodic)[None])
 
-    sin_cos_true_chi = batch['chi_angles_sin_cos'][:, None, ...]
+    sin_cos_true_chi = batch['chi_angles_sin_cos'].unsqueeze(axis=1) # [:, None, ...]
 
     # This is -1 if chi is pi-periodic and +1 if it's 2pi-periodic
     shifted_mask = (1 - 2 * chi_pi_periodic)[..., None]
@@ -913,7 +919,7 @@ def l2_normalize(x, axis=-1, epsilon=1e-12):
     return x / paddle.sqrt(
         paddle.maximum(
             paddle.sum(paddle.square(x), axis=axis, keepdim=True),
-            paddle.to_tensor([epsilon], dtype='float32')))
+            paddle.full(shape=[1], fill_value=epsilon, dtype='float32')))
 
 
 class MultiRigidSidechain(nn.Layer):
@@ -925,9 +931,11 @@ def __init__(self, channel_num, config, global_config):
         self.config = config
         self.global_config = global_config
 
+        Linear = paddle.incubate.nn.FusedLinear if self.global_config.fuse_linear else paddle.nn.Linear
+
         c = self.config.num_channel
-        self.input_projection = nn.Linear(channel_num['seq_channel'], c)
-        self.input_projection_1 = nn.Linear(channel_num['seq_channel'], c)
+        self.input_projection = Linear(channel_num['seq_channel'], c)
+        self.input_projection_1 = Linear(channel_num['seq_channel'], c)
 
         for i in range(self.config.num_residual_block):
             l1, l2 = 'resblock1', 'resblock2'
@@ -940,12 +948,12 @@ def __init__(self, channel_num, config, global_config):
             else:
                 init_w_2 = nn.initializer.XavierUniform()
 
-            setattr(self, l1, nn.Linear(
+            setattr(self, l1, Linear(
                 c, c, weight_attr=paddle.ParamAttr(initializer=init_w_1)))
-            setattr(self, l2, nn.Linear(
+            setattr(self, l2, Linear(
                 c, c, weight_attr=paddle.ParamAttr(initializer=init_w_2)))
 
-        self.unnormalized_angles = nn.Linear(c, 14)
+        self.unnormalized_angles = Linear(c, 14)
 
     def forward(self, affine, single_act, init_single_act, aatype):
         single_act = self.input_projection(nn.functional.relu(single_act))
@@ -997,4 +1005,4 @@ def forward(self, affine, single_act, init_single_act, aatype):
         #     'frames': all_frames_to_global,  # (B, N, 8, 3, 3)
         # })
 
-        return outputs
+        return outputs