test_evoformer.py 10.6 KB
Newer Older
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Copyright 2021 AlQuraishi Laboratory
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import numpy as np
import unittest
18
19
20
21
22
23
24
25
26
from openfold.model.evoformer import (
    MSATransition,
    EvoformerStack,
    ExtraMSAStack,
)
from openfold.utils.tensor_utils import tree_map
import tests.compare_utils as compare_utils
from tests.config import consts

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
27
if compare_utils.alphafold_is_installed():
28
29
30
    alphafold = compare_utils.import_alphafold()
    import jax
    import haiku as hk
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
31
32
33


class TestEvoformerStack(unittest.TestCase):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
34
    def test_shape(self):
35
36
37
38
39
        batch_size = consts.batch_size
        n_seq = consts.n_seq
        n_res = consts.n_res
        c_m = consts.c_m
        c_z = consts.c_z
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
40
41
42
43
        c_hidden_msa_att = 12
        c_hidden_opm = 17
        c_hidden_mul = 19
        c_hidden_pair_att = 14
44
        c_s = consts.c_s
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
        no_heads_msa = 3
        no_heads_pair = 7
        no_blocks = 2
        transition_n = 2
        msa_dropout = 0.15
        pair_stack_dropout = 0.25
        inf = 1e9
        eps = 1e-10

        es = EvoformerStack(
            c_m,
            c_z,
            c_hidden_msa_att,
            c_hidden_opm,
            c_hidden_mul,
            c_hidden_pair_att,
            c_s,
            no_heads_msa,
            no_heads_pair,
            no_blocks,
            transition_n,
            msa_dropout,
            pair_stack_dropout,
            blocks_per_ckpt=None,
69
            no_column_attention=False,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
70
71
72
73
            inf=inf,
            eps=eps,
        ).eval()

74
        m = torch.rand((batch_size, n_seq, n_res, c_m))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
75
        z = torch.rand((batch_size, n_res, n_res, c_z))
76
        msa_mask = torch.randint(0, 2, size=(batch_size, n_seq, n_res))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
77
78
79
80
81
        pair_mask = torch.randint(0, 2, size=(batch_size, n_res, n_res))

        shape_m_before = m.shape
        shape_z_before = z.shape

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
82
83
84
        m, z, s = es(
            m, z, chunk_size=4, msa_mask=msa_mask, pair_mask=pair_mask
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
85
86
87
88
89

        self.assertTrue(m.shape == shape_m_before)
        self.assertTrue(z.shape == shape_z_before)
        self.assertTrue(s.shape == (batch_size, n_res, c_s))

90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
    def test_shape_without_column_attention(self):
        batch_size = consts.batch_size
        n_seq = consts.n_seq
        n_res = consts.n_res
        c_m = consts.c_m
        c_z = consts.c_z
        c_hidden_msa_att = 12
        c_hidden_opm = 17
        c_hidden_mul = 19
        c_hidden_pair_att = 14
        c_s = consts.c_s
        no_heads_msa = 3
        no_heads_pair = 7
        no_blocks = 2
        transition_n = 2
        msa_dropout = 0.15
        pair_stack_dropout = 0.25
        inf = 1e9
        eps = 1e-10

        es = EvoformerStack(
            c_m,
            c_z,
            c_hidden_msa_att,
            c_hidden_opm,
            c_hidden_mul,
            c_hidden_pair_att,
            c_s,
            no_heads_msa,
            no_heads_pair,
            no_blocks,
            transition_n,
            msa_dropout,
            pair_stack_dropout,
            blocks_per_ckpt=None,
            no_column_attention=True,
            inf=inf,
            eps=eps,
        ).eval()

        m_init = torch.rand((batch_size, n_seq, n_res, c_m))
        z_init = torch.rand((batch_size, n_res, n_res, c_z))
        msa_mask = torch.randint(0, 2, size=(batch_size, n_seq, n_res))
        pair_mask = torch.randint(0, 2, size=(batch_size, n_res, n_res))

        shape_m_before = m_init.shape
        shape_z_before = z_init.shape

        m, z, s = es(
            m_init, z_init, chunk_size=4, msa_mask=msa_mask, pair_mask=pair_mask
        )

        self.assertTrue(m.shape == shape_m_before)
        self.assertTrue(z.shape == shape_z_before)
        self.assertTrue(s.shape == (batch_size, n_res, c_s))

146
147
148
149
150
151
    @compare_utils.skip_unless_alphafold_installed()
    def test_compare(self):
        def run_ei(activations, masks):
            config = compare_utils.get_alphafold_config()
            c_e = config.model.embeddings_and_evoformer.evoformer
            ei = alphafold.model.modules.EvoformerIteration(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
152
153
                c_e, config.model.global_config, is_extra_msa=False
            )
154
            return ei(activations, masks, is_training=False)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
155

156
        f = hk.transform(run_ei)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
157

158
159
        n_res = consts.n_res
        n_seq = consts.n_seq
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
160

161
        activations = {
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
162
163
            "msa": np.random.rand(n_seq, n_res, consts.c_m).astype(np.float32),
            "pair": np.random.rand(n_res, n_res, consts.c_z).astype(np.float32),
164
        }
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
165

166
        masks = {
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
167
168
            "msa": np.random.randint(0, 2, (n_seq, n_res)).astype(np.float32),
            "pair": np.random.randint(0, 2, (n_res, n_res)).astype(np.float32),
169
        }
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
170

171
172
173
174
        params = compare_utils.fetch_alphafold_module_weights(
            "alphafold/alphafold_iteration/evoformer/evoformer_iteration"
        )
        params = tree_map(lambda n: n[0], params, jax.numpy.DeviceArray)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
175

176
        key = jax.random.PRNGKey(42)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
177
        out_gt = f.apply(params, key, activations, masks)
178
179
180
        jax.tree_map(lambda x: x.block_until_ready(), out_gt)
        out_gt_msa = torch.as_tensor(np.array(out_gt["msa"]))
        out_gt_pair = torch.as_tensor(np.array(out_gt["pair"]))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
181

182
183
        model = compare_utils.get_global_pretrained_openfold()
        out_repro_msa, out_repro_pair = model.evoformer.blocks[0](
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
184
185
186
            torch.as_tensor(activations["msa"]).cuda(),
            torch.as_tensor(activations["pair"]).cuda(),
            torch.as_tensor(masks["msa"]).cuda(),
187
            torch.as_tensor(masks["pair"]).cuda(),
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
188
            chunk_size=4,
189
            _mask_trans=False,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
190
            inplace_safe=False,
191
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
192

193
194
195
        out_repro_msa = out_repro_msa.cpu()
        out_repro_pair = out_repro_pair.cpu()

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
        self.assertTrue(torch.mean(torch.abs(out_repro_msa - out_gt_msa)) < consts.eps)
        self.assertTrue(torch.max(torch.abs(out_repro_pair - out_gt_pair)) < consts.eps)

        # Inplace version
        out_repro_msa, out_repro_pair = model.evoformer.blocks[0](
            torch.as_tensor(activations["msa"]).cuda(),
            torch.as_tensor(activations["pair"]).cuda(),
            torch.as_tensor(masks["msa"]).cuda(),
            torch.as_tensor(masks["pair"]).cuda(),
            chunk_size=4,
            _mask_trans=False,
            inplace_safe=True,
        )

        out_repro_msa = out_repro_msa.cpu()
        out_repro_pair = out_repro_pair.cpu()

        self.assertTrue(torch.mean(torch.abs(out_repro_msa - out_gt_msa)) < consts.eps)
        self.assertTrue(torch.max(torch.abs(out_repro_pair - out_gt_pair)) < consts.eps)
215

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
216
217

class TestExtraMSAStack(unittest.TestCase):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
218
    def test_shape(self):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
        batch_size = 2
        s_t = 23
        n_res = 5
        c_m = 7
        c_z = 11
        c_hidden_msa_att = 12
        c_hidden_opm = 17
        c_hidden_mul = 19
        c_hidden_tri_att = 16
        no_heads_msa = 3
        no_heads_pair = 8
        no_blocks = 2
        transition_n = 5
        msa_dropout = 0.15
        pair_stack_dropout = 0.25
        inf = 1e9
        eps = 1e-10

        es = ExtraMSAStack(
            c_m,
            c_z,
            c_hidden_msa_att,
            c_hidden_opm,
            c_hidden_mul,
            c_hidden_tri_att,
            no_heads_msa,
            no_heads_pair,
            no_blocks,
            transition_n,
            msa_dropout,
            pair_stack_dropout,
Gustaf Ahdritz's avatar
Fixes  
Gustaf Ahdritz committed
250
            ckpt=False,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
251
252
            inf=inf,
            eps=eps,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
253
        ).eval().cuda()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
254

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
255
256
        m = torch.rand((batch_size, s_t, n_res, c_m), device="cuda")
        z = torch.rand((batch_size, n_res, n_res, c_z), device="cuda")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
257
258
259
260
261
262
263
264
        msa_mask = torch.randint(
            0,
            2,
            size=(
                batch_size,
                s_t,
                n_res,
            ),
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
265
            device="cuda",
266
        ).float()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
267
268
269
270
271
272
273
274
        pair_mask = torch.randint(
            0,
            2,
            size=(
                batch_size,
                n_res,
                n_res,
            ),
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
275
            device="cuda",
276
        ).float()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
277
278
279

        shape_z_before = z.shape

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
280
        z = es(m, z, chunk_size=4, msa_mask=msa_mask, pair_mask=pair_mask)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
281
282
283
284
285

        self.assertTrue(z.shape == shape_z_before)


class TestMSATransition(unittest.TestCase):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
286
    def test_shape(self):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
287
288
289
290
291
292
        batch_size = 2
        s_t = 3
        n_r = 5
        c_m = 7
        n = 11

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
293
        mt = MSATransition(c_m, n)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
294
295
296
297

        m = torch.rand((batch_size, s_t, n_r, c_m))

        shape_before = m.shape
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
298
        m = mt(m, chunk_size=4)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
299
300
301
302
        shape_after = m.shape

        self.assertTrue(shape_before == shape_after)

303
304
305
306
307
308
    @compare_utils.skip_unless_alphafold_installed()
    def test_compare(self):
        def run_msa_transition(msa_act, msa_mask):
            config = compare_utils.get_alphafold_config()
            c_e = config.model.embeddings_and_evoformer.evoformer
            msa_trans = alphafold.model.modules.Transition(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
309
                c_e.msa_transition,
310
                config.model.global_config,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
311
                name="msa_transition",
312
313
314
            )
            act = msa_trans(act=msa_act, mask=msa_mask)
            return act
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
315

316
        f = hk.transform(run_msa_transition)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
317

318
319
        n_res = consts.n_res
        n_seq = consts.n_seq
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
320

321
        msa_act = np.random.rand(n_seq, n_res, consts.c_m).astype(np.float32)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
322
323
324
325
        msa_mask = np.ones((n_seq, n_res)).astype(
            np.float32
        )  # no mask here either

326
327
        # Fetch pretrained parameters (but only from one block)]
        params = compare_utils.fetch_alphafold_module_weights(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
328
329
            "alphafold/alphafold_iteration/evoformer/evoformer_iteration/"
            + "msa_transition"
330
331
        )
        params = tree_map(lambda n: n[0], params, jax.numpy.DeviceArray)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
332
333

        out_gt = f.apply(params, None, msa_act, msa_mask).block_until_ready()
334
        out_gt = torch.as_tensor(np.array(out_gt))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
335

336
        model = compare_utils.get_global_pretrained_openfold()
Gustaf Ahdritz's avatar
Fixes  
Gustaf Ahdritz committed
337
        
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
338
        out_repro = (
Gustaf Ahdritz's avatar
Fixes  
Gustaf Ahdritz committed
339
            model.evoformer.blocks[0].core.msa_transition(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
340
341
342
343
344
345
                torch.as_tensor(msa_act, dtype=torch.float32).cuda(),
                mask=torch.as_tensor(msa_mask, dtype=torch.float32).cuda(),
            )
            .cpu()
        )

Gustaf Ahdritz's avatar
Fixes  
Gustaf Ahdritz committed
346
        self.assertTrue(torch.max(torch.abs(out_gt - out_repro)) < consts.eps)
347

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
348
349
350

if __name__ == "__main__":
    unittest.main()