A method, system, device and storage medium for optimizing protein structure prediction

Abstract

The present invention provides a method, system, device and storage medium for optimizing protein structure prediction, the method comprising: modifying and improving the openfold code to achieve preliminary running through of the second-generation deep learning neural network structure alphafold2; segmenting the data in openfold for preliminary video memory optimization; performing user profile analysis on the second-generation deep learning neural network structure alphafold2 to determine the module with the largest current video memory consumption; and segmenting the data of the module with the largest current video memory consumption for video memory optimization. The present invention can quickly implement the training task of alphafold2, realize the support for long sequence protein training, solve the strong requirements of training on hardware resources, and improve the training performance.

Description

Technical Field

[0001] The present invention relates to the field of deep learning, and more specifically, to a method, system, device and storage medium for optimizing protein structure prediction. Background Art

[0002] With the development of deep learning, its application is becoming more and more extensive. More and more fields are gradually taking deep learning as the key development direction in the future, including the biopharmaceutical industry. In the field of structural biology, protein structure prediction has always been a highly concerned issue. Traditionally, protein structure prediction requires multiple experiments with cryo-electron microscopy. The time spent on each protein structure prediction is calculated in years, which consumes too much manpower and material resources. However, in recent years, with the rapid development of deep learning, the accuracy of protein structures predicted by deep learning has once exceeded the experimental results, saving researchers a lot of time and energy.

[0003] As a leading algorithm for protein structure prediction, alphafold2 has attracted widespread attention since it was open-sourced. However, for researchers, there are many problems in using alphafold2 for further research or commercial purposes. First, alphafold2 only open-sources the inference code and cannot retrain or fine-tune the model. Second, model training requires a lot of video memory support. During finetuning, the video memory usage of a single attention logits reaches 20.25GB, which cannot support the training of long sequence proteins. The entire training process requires huge hardware resource support.

[0004] Openfold is an open source project that uses pytorch to reproduce alphafold2. Its reproduction is relatively complete compared to other projects, but there are many bugs (vulnerabilities) in order to run the code successfully for training. Fastfold is a current memory optimization method for alphafold2 reasoning. It only optimizes the evoformer module of alphafold2. If you want to apply it to the training of alphafold2 and realize the training of long sequence proteins, a lot of work is still needed. Summary of the invention

[0005] In view of this, the purpose of the embodiments of the present invention is to propose a method, system, computer device and computer-readable storage medium for optimizing protein structure prediction. The present invention helps researchers to quickly implement the training process of alphafold2 and supports training of long sequences. By using this method, not only can the entire training and finetuning of alphafold2 be quickly implemented, but also the time of each iteration step in the training process can be reduced, which provides a convenient idea for scientific researchers to carry out further work based on alphafold2, saving a lot of time, material resources and manpower.

[0006] Based on the above purpose, one aspect of an embodiment of the present invention provides a method for optimizing protein structure prediction, including the following steps: modifying and improving the openfold code to achieve preliminary running of the second-generation deep learning neural network structure alphafold2; splitting the data in openfold for preliminary video memory optimization; performing user profile analysis on the second-generation deep learning neural network structure alphafold2 to determine the module with the largest current video memory consumption; and splitting the data of the module with the largest current video memory consumption for video memory optimization.

[0007] In some embodiments, the modification and improvement of the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2 includes: using a while loop to control openfold to read files step by step, and training openfold with FP32 precision.

[0008] In some embodiments, splitting the data in openfold for preliminary video memory optimization includes: copying the same data into multiple copies, and splitting the data through a scatter function so that each device is allocated a portion of the data for calculation; and in response to the completion of the calculation, merging the data using an aggregation function.

[0009] In some embodiments, the splitting of the data in OpenFold for preliminary video memory optimization includes: modifying the distributed training DDP of PyTorch to the DDP of Apex, and modifying the softmax operator to the softmax provided by Torch.

[0010] Another aspect of an embodiment of the present invention provides a system for optimizing protein structure prediction, including: a modification module, configured to modify and improve the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2; a first video memory module, configured to split the data in openfold for preliminary video memory optimization; an analysis module, configured to perform user profile analysis on the second-generation deep learning neural network structure alphafold2 to determine the module with the largest current video memory consumption; and a second video memory module, configured to split the data of the module with the largest current video memory consumption for video memory optimization.

[0011] In some embodiments, the modification module is configured to: use a while loop to control openfold to read files step by step, and train openfold using FP32 precision.

[0012] In some embodiments, the first video memory module is configured to: copy the same data into multiple copies, and divide the data by a scatter function so that each device is allocated a portion of the data for calculation; and in response to the completion of the calculation, merge the data using an aggregation function.

[0013] In some embodiments, the first video memory module is configured to: modify the distributed training DDP of pytorch to the DDP of apex, and modify the softmax operator to the softmax that comes with torch.

[0014] According to another aspect of an embodiment of the present invention, a computer device is provided, comprising: at least one processor; and a memory, wherein the memory stores computer instructions executable on the processor, and the instructions implement the steps of the above method when executed by the processor.

[0015] According to another aspect of the embodiments of the present invention, a computer-readable storage medium is provided, which stores a computer program that implements the above method steps when executed by a processor.

[0016] The present invention has the following beneficial technical effects: it helps researchers to quickly implement the training process of alphafold2 and supports the training of long sequences. By using this method, not only can the entire training and finetuning of alphafold2 be quickly implemented, but also the time of each iteration step in the training process can be reduced, which provides a convenient way for scientific researchers to carry out further work based on alphafold2, saving a lot of time, material resources and manpower. BRIEF DESCRIPTION OF THE DRAWINGS

[0017] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings in the following description are only some embodiments of the present invention. For ordinary technicians in this field, other embodiments can be obtained based on these drawings without paying creative work.

[0018] Figure 1 A schematic diagram of an embodiment of a method for optimizing protein structure prediction provided by the present invention;

[0019] Figure 2 A schematic diagram of an embodiment of a system for optimizing protein structure prediction provided by the present invention;

[0020] Figure 3 A schematic diagram of the hardware structure of an embodiment of a computer device for optimizing protein structure prediction provided by the present invention;

[0021] Figure 4 A schematic diagram of an embodiment of a computer storage medium for optimizing protein structure prediction provided by the present invention. DETAILED DESCRIPTION

[0022] In order to make the objectives, technical solutions and advantages of the present invention more clearly understood, the embodiments of the present invention are further described in detail below in combination with specific embodiments and with reference to the accompanying drawings.

[0023] It should be noted that all expressions using "first" and "second" in the embodiments of the present invention are for distinguishing two non-identical entities with the same name or non-identical parameters. It can be seen that "first" and "second" are only for the convenience of expression and should not be understood as limitations on the embodiments of the present invention. The subsequent embodiments will not explain this one by one.

[0024] The first aspect of the embodiments of the present invention provides an embodiment of a method for optimizing protein structure prediction. Figure 1 What is shown is a schematic diagram of an embodiment of the method for optimizing protein structure prediction provided by the present invention. Figure 1 As shown, the embodiment of the present invention includes the following steps:

[0025] S1. Modify and improve the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2;

[0026] S2, split the data in openfold for preliminary video memory optimization;

[0027] S3, performing a user profile analysis on the second-generation deep learning neural network structure alphafold2 to determine the module with the largest current video memory consumption; and

[0028] S4. Split the data of the module that currently consumes the most video memory to optimize the video memory.

[0029] As a product of the combination of biological structure and artificial intelligence, alphafold2 has attracted widespread attention from the scientific research community and the biological community as soon as it was open sourced. Since it comes from deepmind, the training of alphafold2 is based on TPU, so its code itself is inherently insufficient in supporting GPUs, and the deep learning framework used is deepmind's self-developed framework jax, so it is not convenient to use. Through code reproduction, it can be found that its training requires a lot of hardware resource support. How to implement the training of alphafold2 and reduce its video memory occupancy and implement the training of long sequence proteins is an urgent problem that many scholars and researchers need to solve. The present invention proposes a method for implementing alphafold2 training and video memory optimization, which helps researchers quickly implement the training process of alphafold2 and supports the training of long sequences. By using this method, not only can the entire training and finetuning of alphafold2 be quickly implemented, but also the time of each iteration in the training process can be reduced, which provides a convenient way for scientific researchers to carry out further work based on alphafold2, saving a lot of time, material resources and manpower.

[0030] Modify and improve the openfold code to achieve a preliminary run through the second-generation deep learning neural network structure alphafold2.

[0031] In some embodiments, the modification and improvement of the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2 includes: using a while loop to control openfold to read files step by step, and training openfold with FP32 precision.

[0032] Openfold has the problem of data not existing during the data reading stage, so it often interrupts. The main reason is that some input cif files have replacement files, and some replacement files have replacement files. Therefore, it is necessary to use a while loop to read the files step by step. For files that still do not exist in the end, use the try except operation to return None (empty). In addition, in the model pre-processing stage, openfold uses fp16 as the training accuracy, but long-term use of fp16 will find that the loss has the problem of NaN, so for openfold training, all FP32 accuracy must be used.

[0033] Split the data in openfold for preliminary video memory optimization.

[0034] In some embodiments, splitting the data in openfold for preliminary video memory optimization includes: copying the same data into multiple copies, and splitting the data through a scatter function so that each device is allocated a portion of the data for calculation; and in response to the completion of the calculation, merging the data using an aggregation function.

[0035] In some embodiments, the splitting of the data in OpenFold for preliminary video memory optimization includes: modifying the distributed training DDP of PyTorch to the DDP of Apex, and modifying the softmax operator to the softmax provided by Torch.

[0036] Fastfold optimizes the video memory of the evoformer module in alphafold2. Its main idea is to divide the data into multiple copies and input them to different devices respectively. After calculation, the output results are merged, that is, the same data is distributed to multiple devices for calculation, so as to achieve the purpose of reducing video memory. The first step to apply it to openfold is to split the data. First, you need to rewrite torch.utils.data.distributed.DistributedSampler to copy the same data into multiple copies to ensure that the devices processing the same data are assigned to the same batch of data. Then, the data is split by scatter (dispersion function), and each device is assigned a part of the data for calculation. After the calculation is completed, gather (aggregation function) is used to merge the data; because fastfold data merging involves communication problems and there is a blockage with pytorch's ddp, we changed pytorch's distributed training ddp to apex's ddp. In addition, in addition to proposing to reduce the cache by splitting the data, fastfold also modified the layernorm and softmax operators, but the new softmax has many problems, so it needs to be modified to torch's built-in softmax when using it.

[0037] The user profile analysis of the second generation deep learning neural network structure alphafold2 is performed to determine the module with the largest video memory consumption. The data of the module with the largest video memory consumption is split for video memory optimization. Through profile analysis, it can be found that the extra_msa_stack part of the alphafold2 structure consumes the most video memory. Using the idea of ​​data splitting, this module is also optimized for video memory.

[0038] Through optimization, the video memory usage has been reduced from 28GB to 13.6GB at 256 sequences, a reduction of more than 50%, and the training performance has been improved by 89%. At 384 sequences, the problem of Out of Memory originally occurred, but now the training can be done normally, and the video memory usage is 30.8GB, which is enough to support training on the 40GB A100.

[0039] It should be particularly pointed out that the various steps in the various embodiments of the above-mentioned method for optimizing protein structure prediction can be cross-linked, replaced, added, and deleted. Therefore, these reasonable permutations, combinations, and transformations for the method for optimizing protein structure prediction should also fall within the scope of protection of the present invention, and the scope of protection of the present invention should not be limited to the embodiments.

[0040] Based on the above objectives, the second aspect of the embodiments of the present invention provides a system for optimizing protein structure prediction. Figure 2 As shown, the system 200 includes the following modules: a modification module, configured to modify and improve the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2; a first video memory module, configured to split the data in openfold for preliminary video memory optimization; an analysis module, configured to perform a user profile analysis on the second-generation deep learning neural network structure alphafold2 to determine the module with the largest current video memory consumption; and a second video memory module, configured to split the data of the module with the largest current video memory consumption for video memory optimization.

[0041] In some embodiments, the modification module is configured to: use a while loop to control openfold to read files step by step, and train openfold using FP32 precision.

[0042] In some embodiments, the first video memory module is configured to: copy the same data into multiple copies, and divide the data by a scatter function so that each device is allocated a portion of the data for calculation; and in response to the completion of the calculation, merge the data using an aggregation function.

[0043] In some embodiments, the first video memory module is configured to: modify the distributed training DDP of pytorch to the DDP of apex, and modify the softmax operator to the softmax that comes with torch.

[0044] Based on the above purpose, the third aspect of an embodiment of the present invention proposes a computer device, including: at least one processor; and a memory, the memory storing computer instructions that can be run on the processor, the instructions are executed by the processor to implement the following steps: S1, modify and improve the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2; S2, split the data in openfold for preliminary video memory optimization; S3, perform user profile analysis on the second-generation deep learning neural network structure alphafold2 to determine the module with the largest current video memory consumption; and S4, split the data of the module with the largest current video memory consumption for video memory optimization.

[0045] In some embodiments, the modification and improvement of the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2 includes: using a while loop to control openfold to read files step by step, and training openfold with FP32 precision.

[0046] In some embodiments, splitting the data in openfold for preliminary video memory optimization includes: copying the same data into multiple copies, and splitting the data through a scatter function so that each device is allocated a portion of the data for calculation; and in response to the completion of the calculation, merging the data using an aggregation function.

[0047] In some embodiments, the splitting of the data in OpenFold for preliminary video memory optimization includes: modifying the distributed training DDP of PyTorch to the DDP of Apex, and modifying the softmax operator to the softmax provided by Torch.

[0048] like Figure 3 , which is a schematic diagram of the hardware structure of an embodiment of the computer device for optimizing protein structure prediction provided by the present invention.

[0049] As Figure 3 Taking the device shown as an example, the device includes a processor 301 and a memory 302.

[0050] The processor 301 and the memory 302 may be connected via a bus or other means. Figure 3 The example of connecting through bus is taken in the following.

[0051] The memory 302 is a non-volatile computer-readable storage medium that can be used to store non-volatile software programs, non-volatile computer executable programs and modules, such as program instructions / modules corresponding to the method for optimizing protein structure prediction in the embodiment of the present application. The processor 301 executes various functional applications and data processing of the server by running the non-volatile software programs, instructions and modules stored in the memory 302, that is, implementing the method for optimizing protein structure prediction.

[0052] The memory 302 may include a program storage area and a data storage area, wherein the program storage area may store an operating system, an application required for at least one function; the data storage area may store data created according to the use of the method for optimizing protein structure prediction, etc. In addition, the memory 302 may include a high-speed random access memory, and may also include a non-volatile memory, such as at least one disk storage device, a flash memory device, or other non-volatile solid-state storage device. In some embodiments, the memory 302 may optionally include a memory remotely arranged relative to the processor 301, and these remote memories may be connected to the local module via a network. Examples of the above-mentioned network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0053] Computer instructions 303 corresponding to one or more methods for optimizing protein structure prediction are stored in the memory 302 , and when executed by the processor 301 , the method for optimizing protein structure prediction in any of the above method embodiments is executed.

[0054] Any embodiment of the computer device that executes the above-mentioned method for optimizing protein structure prediction can achieve the same or similar effect as the corresponding any of the above-mentioned method embodiments.

[0055] The present invention also provides a computer-readable storage medium, which stores a computer program for executing the method for optimizing protein structure prediction when executed by a processor.

[0056] like Figure 4 FIG. 1 is a schematic diagram of an embodiment of the computer storage medium for optimizing protein structure prediction provided by the present invention. Figure 4 Taking the computer storage medium shown as an example, the computer readable storage medium 401 stores a computer program 402 that performs the above method when executed by a processor.

[0057] Finally, it should be noted that a person skilled in the art can understand that all or part of the processes in the above-mentioned embodiment method can be implemented by instructing the relevant hardware through a computer program, and the program of the method for optimizing protein structure prediction can be stored in a computer-readable storage medium, and when the program is executed, it can include the processes of the embodiments of the above-mentioned methods. Among them, the storage medium of the program can be a disk, an optical disk, a read-only storage memory (ROM) or a random access memory (RAM), etc. The above-mentioned computer program embodiment can achieve the same or similar effect as the corresponding above-mentioned arbitrary method embodiment.

[0058] The above are exemplary embodiments disclosed in the present invention, but it should be noted that various changes and modifications may be made without departing from the scope disclosed in the embodiments of the present invention as defined in the claims. The functions, steps and / or actions of the method claims according to the disclosed embodiments described herein do not need to be performed in any particular order. In addition, although the elements disclosed in the embodiments of the present invention may be described or required in individual form, they may also be understood as multiple unless explicitly limited to the singular.

[0059] It should be understood that, as used herein, the singular forms "a", "an" are intended to include the plural forms as well, unless the context clearly supports an exception. It should also be understood that, as used herein, "and / or" refers to any and all possible combinations including one or more of the associated listed items.

[0060] The serial numbers of the embodiments disclosed in the above embodiments of the present invention are only for description and do not represent the advantages or disadvantages of the embodiments.

[0061] A person skilled in the art will understand that all or part of the steps to implement the above embodiments may be accomplished by hardware or by instructing related hardware through a program, and the program may be stored in a computer-readable storage medium, and the above-mentioned storage medium may be a read-only memory, a disk or an optical disk, etc.

[0062] A person skilled in the art should understand that the discussion of any of the above embodiments is only exemplary and is not intended to imply that the scope of the disclosure of the embodiments of the present invention (including the claims) is limited to these examples; under the concept of the embodiments of the present invention, the technical features in the above embodiments or different embodiments can also be combined, and there are many other changes in different aspects of the embodiments of the present invention as above, which are not provided in detail for the sake of simplicity. Therefore, any omissions, modifications, equivalent substitutions, improvements, etc. made within the spirit and principles of the embodiments of the present invention should be included in the protection scope of the embodiments of the present invention.

Claims

1. A method for optimizing protein structure prediction, characterized in that: The steps include: Modify and improve the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2; Split the data in openfold for preliminary video memory optimization; Perform user profile analysis on the second-generation deep learning neural network architecture, alphafold2, to determine the modules that currently consume the most video memory; as well as Split the data of the module that currently consumes the most video memory to optimize the video memory; The modification and improvement of the openfold code to achieve the preliminary run-through of the second-generation deep learning neural network structure alphafold2 includes: using a while loop to control openfold to read files step by step, and using FP32 precision to train openfold; The splitting of the data in the openfold for preliminary video memory optimization includes: copying the same data into multiple copies, and splitting the data by a scatter function so that each device is allocated a portion of the data for calculation; and in response to the completion of the calculation, merging the data by using an aggregation function; Change the distributed training DDP of pytorch to the DDP of apex, and change the softmax operator to the softmax that comes with torch.

2. A system for optimizing protein structure prediction, characterized in that: include: Modification module, configured to modify and improve the openfold code to achieve a preliminary run-through of the second-generation deep learning neural network structure alphafold2; The first video memory module is configured to split the data in OpenFold for preliminary video memory optimization; An analysis module configured to perform a user profile analysis on the second-generation deep learning neural network structure alphafold2 to determine a module with the largest current video memory consumption; as well as A second video memory module is configured to perform data segmentation on the module with the largest video memory consumption to optimize the video memory; The modification module is configured to: use a while loop to control openfold to read files step by step, and train openfold using FP32 precision; The first video memory module is configured to: copy the same data into multiple copies, and divide the data by a scatter function so that each device is allocated a portion of the data for calculation; and in response to the completion of the calculation, merge the data using an aggregation function; The first video memory module is configured to: modify the distributed training DDP of pytorch to the DDP of apex, and modify the softmax operator to the softmax that comes with torch.

3. A computer device, characterized in that: include: at least one processor; as well as A memory storing computer instructions executable on the processor, wherein the instructions implement the steps of the method according to claim 1 when executed by the processor.

4. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the steps of the method according to claim 1 are implemented.

Images