run_pretrained_openfold.py 15.8 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
# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
# 
# 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.

16
import argparse
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
17
from datetime import date
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
18
import gc
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
19
import logging
20
import numpy as np
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
21
import os
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
22

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
23
import pickle
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
24
25
26
from pytorch_lightning.utilities.deepspeed import (
    convert_zero_checkpoint_to_fp32_state_dict
)
27
28
import random
import sys
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
29
30
31
import time
import torch

32
from openfold.config import model_config
33
from openfold.data import templates, feature_pipeline, data_pipeline
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
34
from openfold.model.model import AlphaFold
35
from openfold.model.torchscript import script_preset_
36
from openfold.np import residue_constants, protein
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
37
38
import openfold.np.relax.relax as relax
from openfold.utils.import_weights import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
39
40
    import_jax_weights_,
)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
41
from openfold.utils.tensor_utils import (
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
42
43
44
    tensor_tree_map,
)

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
45
from scripts.utils import add_data_args
46

47

48
49
50
51
52
logging.basicConfig()
logger = logging.getLogger(__file__)
logger.setLevel(level=logging.INFO)


Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
53
54
55
56
57
58
59
60
def precompute_alignments(tags, seqs, alignment_dir, args):
    for tag, seq in zip(tags, seqs):
        tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
        with open(tmp_fasta_path, "w") as fp:
            fp.write(f">{tag}\n{seq}")

        local_alignment_dir = os.path.join(alignment_dir, tag)
        if(args.use_precomputed_alignments is None):
61
            logger.info(f"Generating alignments for {tag}...") 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
            if not os.path.exists(local_alignment_dir):
                os.makedirs(local_alignment_dir)
            
            alignment_runner = data_pipeline.AlignmentRunner(
                jackhmmer_binary_path=args.jackhmmer_binary_path,
                hhblits_binary_path=args.hhblits_binary_path,
                hhsearch_binary_path=args.hhsearch_binary_path,
                uniref90_database_path=args.uniref90_database_path,
                mgnify_database_path=args.mgnify_database_path,
                bfd_database_path=args.bfd_database_path,
                uniclust30_database_path=args.uniclust30_database_path,
                pdb70_database_path=args.pdb70_database_path,
                use_small_bfd=use_small_bfd,
                no_cpus=args.cpus,
            )
            alignment_runner.run(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
78
                fasta_path, local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
79
80
81
82
83
84
85
86
87
88
89
90
91
92
            )

        # Remove temporary FASTA file
        os.remove(tmp_fasta_path)


def run_model(model, batch, tag, args):
    with torch.no_grad():
        batch = {
            k:torch.as_tensor(v, device=args.model_device) 
            for k,v in batch.items()
        }
 
        # Disable templates if there aren't any in the batch
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
93
        model.config.template.enabled = model.config.template.enabled and any([
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
94
95
96
            "template_" in k for k in batch
        ])

97
        logger.info(f"Running inference for {tag}...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
98
99
        t = time.perf_counter()
        out = model(batch)
100
        logger.info(f"Inference time: {time.perf_counter() - t}")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    
    return out


def prep_output(out, batch, feature_dict, feature_processor, args):
    plddt = out["plddt"]
    mean_plddt = np.mean(plddt)
    
    plddt_b_factors = np.repeat(
        plddt[..., None], residue_constants.atom_type_num, axis=-1
    )

    # Prep protein metadata
    template_domain_names = []
    template_chain_index = None
116
    if(feature_processor.config.common.use_templates and "template_domain_names" in feature_dict):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        template_domain_names = [
            t.decode("utf-8") for t in feature_dict["template_domain_names"]
        ]

        # This works because templates are not shuffled during inference
        template_domain_names = template_domain_names[
            :feature_processor.config.predict.max_templates
        ]

        if("template_chain_index" in feature_dict):
            template_chain_index = feature_dict["template_chain_index"]
            template_chain_index = template_chain_index[
                :feature_processor.config.predict.max_templates
            ]

    no_recycling = feature_processor.config.common.max_recycling_iters
    remark = ', '.join([
        f"no_recycling={no_recycling}",
        f"max_templates={feature_processor.config.predict.max_templates}",
        f"config_preset={args.model_name}",
    ])

    # For multi-chain FASTAs
    ri = feature_dict["residue_index"]
    chain_index = (ri - np.arange(ri.shape[0])) / args.multimer_ri_gap
    chain_index = chain_index.astype(np.int64)
    cur_chain = 0
    prev_chain_max = 0
    for i, c in enumerate(chain_index):
        if(c != cur_chain):
            cur_chain = c
            prev_chain_max = i + cur_chain * args.multimer_ri_gap

        batch["residue_index"][i] -= prev_chain_max

    unrelaxed_protein = protein.from_prediction(
        features=batch,
        result=out,
        b_factors=plddt_b_factors,
        chain_index=chain_index,
        remark=remark,
        parents=template_domain_names,
        parents_chain_index=template_chain_index,
    )

    return unrelaxed_protein


165
def main(args):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
166
167
168
169
    # Create the output directory
    os.makedirs(args.output_dir, exist_ok=True)

    # Prep the model
170
    config = model_config(args.model_name)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
171
    
172
173
    logger.info(f"Using config preset {args.model_name}...")

174
    model = AlphaFold(config)
175
    model = model.eval()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
176
177
178
179
180

    if(args.jax_param_path):
        import_jax_weights_(
            model, args.jax_param_path, version=args.model_name
        )
181
182
183
        logger.info(
            f"Successfully loaded JAX parameters at {args.jax_param_path}..."
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
184
185
    elif(args.openfold_checkpoint_path):
        if(os.path.isdir(args.openfold_checkpoint_path)):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
186
            # A DeepSpeed checkpoint
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
            checkpoint_basename = os.path.splitext(
                os.path.basename(
                    os.path.normpath(args.openfold_checkpoint_path)
                )
            )[0]
            ckpt_path = os.path.join(
                args.output_dir,
                checkpoint_basename + ".pt",
            ) 

            if(not os.path.isfile(ckpt_path)):
                convert_zero_checkpoint_to_fp32_state_dict(
                    args.openfold_checkpoint_path,
                    ckpt_path,
                )
202
203
204

            d = torch.load(ckpt_path)
            model.load_state_dict(d["ema"]["params"])
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
205
        else:
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
206
207
            # A checkpoint from the public release, which only contains EMA
            # params
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
208
            ckpt_path = args.openfold_checkpoint_path
209
            d = torch.load(ckpt_path)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
210

211
212
213
            if("ema" in d):
                # The public weights have had this done to them already
                d = d["ema"]["params"]
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
214

215
            model.load_state_dict(d)
216
217
218
219
       
        logger.info(
            f"Loaded OpenFold parameters at {args.openfold_checkpoint_path}..."
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
220
221
222
223
224
225
    else:
        raise ValueError(
            "At least one of jax_param_path or openfold_checkpoint_path must "
            "be specified."
        )

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
226
    model = model.to(args.model_device)
227

228
229
230
    template_featurizer = templates.TemplateHitFeaturizer(
        mmcif_dir=args.template_mmcif_dir,
        max_template_date=args.max_template_date,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
231
        max_hits=config.data.predict.max_templates,
232
        kalign_binary_path=args.kalign_binary_path,
233
        release_dates_path=args.release_dates_path,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
234
235
        obsolete_pdbs_path=args.obsolete_pdbs_path
    )
236

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
237
238
    use_small_bfd=(args.bfd_database_path is None)

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
239
    data_processor = data_pipeline.DataPipeline(
240
241
242
243
        template_featurizer=template_featurizer,
    )

    output_dir_base = args.output_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
244
    random_seed = args.data_random_seed
245
246
    if random_seed is None:
        random_seed = random.randrange(sys.maxsize)
247
    feature_processor = feature_pipeline.FeaturePipeline(config.data)
248
249
    if not os.path.exists(output_dir_base):
        os.makedirs(output_dir_base)
Gustaf's avatar
Gustaf committed
250
    if(args.use_precomputed_alignments is None):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
251
        alignment_dir = os.path.join(output_dir_base, "alignments")
Gustaf's avatar
Gustaf committed
252
253
    else:
        alignment_dir = args.use_precomputed_alignments
254
        logger.info(f"Using precomputed alignments at {alignment_dir}...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
255

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
256
    prediction_dir = os.path.join(args.output_dir, "predictions") 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
257
258
    os.makedirs(prediction_dir, exist_ok=True)

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
259
260
261
262
    for fasta_file in os.listdir(args.fasta_dir):
        # Gather input sequences
        with open(os.path.join(args.fasta_dir, fasta_file), "r") as fp:
            data = fp.read()
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
263

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
264
265
266
267
268
        lines = [
            l.replace('\n', '') 
            for prot in data.split('>') for l in prot.strip().split('\n', 1)
        ][1:]
        tags, seqs = lines[::2], lines[1::2]
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
269

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
270
        tags = [t.split()[0] for t in tags]
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
271
        # assert len(tags) == len(set(tags)), "All FASTA tags must be unique"
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
272
        tag = '-'.join(tags)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
273

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
274
275
276
277
278
279
280
281
282
283
284
285
        output_name = f'{tag}_{args.model_name}'
        if(args.output_postfix is not None):
            output_name = f'{output_name}_{args.output_postfix}'

        # Save the unrelaxed PDB.
        unrelaxed_output_path = os.path.join(
            prediction_dir, f'{output_name}_unrelaxed.pdb'
        )

        if(os.path.exists(unrelaxed_output_path)):
            continue

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
286
        precompute_alignments(tags, seqs, alignment_dir, args)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
287

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
288
289
290
291
292
293
294
295
296
        tmp_fasta_path = os.path.join(args.output_dir, f"tmp_{os.getpid()}.fasta")
        if(len(seqs) == 1):
            seq = seqs[0]
            with open(tmp_fasta_path, "w") as fp:
                fp.write(f">{tag}\n{seq}")

            local_alignment_dir = os.path.join(alignment_dir, tag)
            feature_dict = data_processor.process_fasta(
                fasta_path=tmp_fasta_path, alignment_dir=local_alignment_dir
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
297
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
298
299
300
301
302
303
304
        else:
            with open(tmp_fasta_path, "w") as fp:
                fp.write(
                    '\n'.join([f">{tag}\n{seq}" for tag, seq in zip(tags, seqs)])
                )
            feature_dict = data_processor.process_multiseq_fasta(
                fasta_path=tmp_fasta_path, super_alignment_dir=alignment_dir, 
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
305
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
306

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
307
        # Remove temporary FASTA file
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
308
        os.remove(tmp_fasta_path)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
309

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
310
311
312
        processed_feature_dict = feature_processor.process_features(
            feature_dict, mode='predict',
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
313

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
314
        batch = processed_feature_dict
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
315
316
        out = run_model(model, batch, tag, args)

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
317
318
319
        # Toss out the recycling dimensions --- we don't need them anymore
        batch = tensor_tree_map(lambda x: np.array(x[..., -1].cpu()), batch)
        out = tensor_tree_map(lambda x: np.array(x.cpu()), out)
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
320
       
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
321
322
        unrelaxed_protein = prep_output(
            out, batch, feature_dict, feature_processor, args
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
323
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
324

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
325
326
        output_name = f'{tag}_{args.model_name}'
        if(args.output_postfix is not None):
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
327
            output_name = f'{output_name}_{args.output_postfix}_{tag_postfix}'
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
328

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
329
330
        # Save the unrelaxed PDB.
        unrelaxed_output_path = os.path.join(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
331
            prediction_dir, f'{output_name}_unrelaxed.pdb'
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
332
        )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
333
334
        with open(unrelaxed_output_path, 'w') as fp:
            fp.write(protein.to_pdb(unrelaxed_protein))
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
335

336
337
        logger.info(f"Output written to {unrelaxed_output_path}...")

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
338
339
340
341
342
343
344
        if(not args.skip_relaxation):
            amber_relaxer = relax.AmberRelaxation(
                use_gpu=(args.model_device != "cpu"),
                **config.relax,
            )
            
            # Relax the prediction.
345
            logger.info(f"Running relaxation on {unrelaxed_output_path}...")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
346
347
348
349
350
351
352
            t = time.perf_counter()
            visible_devices = os.getenv("CUDA_VISIBLE_DEVICES", default="")
            if("cuda" in args.model_device):
                device_no = args.model_device.split(":")[-1]
                os.environ["CUDA_VISIBLE_DEVICES"] = device_no
            relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
            os.environ["CUDA_VISIBLE_DEVICES"] = visible_devices
353
            logger.info(f"Relaxation time: {time.perf_counter() - t}")
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
354
355
356
            
            # Save the relaxed PDB.
            relaxed_output_path = os.path.join(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
357
                prediction_dir, f'{output_name}_relaxed.pdb'
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
358
            )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
359
360
            with open(relaxed_output_path, 'w') as fp:
                fp.write(relaxed_pdb_str)
361

362
363
            logger.info(f"Relaxed output written to {relaxed_output_path}...")

364
365
        if(args.save_outputs):
            output_dict_path = os.path.join(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
366
                args.output_dir, f'{output_name}_output_dict.pkl'
367
368
369
370
            )
            with open(output_dict_path, "wb") as fp:
                pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)

371
            logger.info(f"Model output written to {output_dict_path}...")
372
373
374

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
375
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
376
377
        "fasta_dir", type=str,
        help="Path to directory containing FASTA files, one sequence per file"
378
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
379
380
381
    parser.add_argument(
        "template_mmcif_dir", type=str,
    )
Gustaf's avatar
Gustaf committed
382
383
384
385
    parser.add_argument(
        "--use_precomputed_alignments", type=str, default=None,
        help="""Path to alignment directory. If provided, alignment computation 
                is skipped and database path arguments are ignored."""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
386
    )
387
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
388
389
        "--output_dir", type=str, default=os.getcwd(),
        help="""Name of the directory in which to output the prediction""",
390
391
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
392
        "--model_device", type=str, default="cpu",
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
393
394
        help="""Name of the device on which to run the model. Any valid torch
             device name is accepted (e.g. "cpu", "cuda:0")"""
395
396
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
397
398
399
        "--model_name", type=str, default="model_1",
        help="""Name of a model config. Choose one of model_{1-5} or 
             model_{1-5}_ptm, as defined on the AlphaFold GitHub."""
400
401
    )
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
402
403
404
405
406
407
408
409
410
        "--jax_param_path", type=str, default=None,
        help="""Path to JAX model parameters. If None, and openfold_checkpoint_path
             is also None, parameters are selected automatically according to 
             the model name from openfold/resources/params"""
    )
    parser.add_argument(
        "--openfold_checkpoint_path", type=str, default=None,
        help="""Path to OpenFold checkpoint. Can be either a DeepSpeed 
             checkpoint directory or a .pt file"""
411
    )
412
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
413
        "--save_outputs", action="store_true", default=False,
414
415
        help="Whether to save all model outputs, including embeddings, etc."
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
416
417
    parser.add_argument(
        "--cpus", type=int, default=4,
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
418
        help="""Number of CPUs with which to run alignment tools"""
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
419
    )
420
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
421
        "--preset", type=str, default='full_dbs',
422
423
        choices=('reduced_dbs', 'full_dbs')
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
424
425
426
427
    parser.add_argument(
        "--output_postfix", type=str, default=None,
        help="""Postfix for output prediction filenames"""
    )
428
    parser.add_argument(
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
429
430
431
432
        "--data_random_seed", type=str, default=None
    )
    parser.add_argument(
        "--skip_relaxation", action="store_true", default=False,
433
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
434
435
436
437
    parser.add_argument(
        "--multimer_ri_gap", type=int, default=200,
        help="""Residue index offset between multiple sequences, if provided"""
    )
Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
438
    add_data_args(parser)
439
440
    args = parser.parse_args()

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
441
442
    if(args.jax_param_path is None and args.openfold_checkpoint_path is None):
        args.jax_param_path = os.path.join(
443
444
445
446
            "openfold", "resources", "params", 
            "params_" + args.model_name + ".npz"
        )

Gustaf Ahdritz's avatar
Gustaf Ahdritz committed
447
448
449
450
451
452
    if(args.model_device == "cpu" and torch.cuda.is_available()):
        logging.warning(
            """The model is being run on CPU. Consider specifying 
            --model_device for better performance"""
        )

453
    main(args)