JointUDL

JointUDL is a deep learning framework for image alignment aiming at cryo-EM two-dimensional single particle images. Aiming to solve rigid image alignment task in noisy situations, we propose a new unsupervised deep learning strategy UDL and further introduce joint unsupervised deep learning method JUDL based on it. The whole program is implemented in Python3 using PyTorch package.

Download   Learn more

Unsupervised Learning Strategies

1.Unsupervised Deep Learning Strategy UDL

UDL is proposed to solve rigid image alignment in noisy situations. It utilizes quantity property of network outputs in regression problem and create pseudo labels for network training. Using such pseudo labels, the original unsolvable unsupervised problem is converted to a solvable pseudo supervised problem. The converted pseudo supervised problem is not equivalent to the original one but serves as an expanded problem which provides a solution space containing the desired solution to the original problem. By properly formulating the difference functions, solutions in solution space would be in a trackable form and the desired solution would be able to be obtained.

2.Joint Unsupervised Deep Learning Method JUDL

JUDL is developed on the basis of pair-wise image alignment method utilizing UDL. Cryo-EM single particle image alignment aims to get a set of images of the same category into the same position. In such situation, similarity of the whole dataset can be utilized for image alignment instead of only pair-wise images. In implementation, when images are not aligned, one batch of images would contain more information than aligned images and thus lead to higher entropy. Therefore, a reconstruction module such as autoencoder would encourage input image to be aligned since aligned images would get lower reconstruction loss. By combining the reconstruction loss with the pair-wise alignment loss, the network would produce more accurate results.

Network Architectures

1.Mask Module m(∙)

Mask prediction module consists of four 2D convolutional layers. It would help network ignore background and focus on particles, and thus learn features containing richer transformation information.

2.Feature Extraction Module fe(∙)

Feature extraction module consists of four residual blocks to learn deep features for image alignment. Unlike entire ResNet or VGG, features are cropped at middle stage to extract features so that there exists enough spatial information for the following feature matching module.

3.Feature Matching Module fm(∙)

Feature matching module consists of one feature matching layer and two 2D convolutional layers. Feature matching layer is essentially a normalized cross-correlation function in which every position of one feature containing all channels is compared to all positions of the other feature to generate correlation map. Siamese structure is utilized to explore more spatial correlation information

4.Regression Module r(∙)

Regression module simply consists of three fully connected layers. After feature matching, two correlation maps are flattened into one dimensional array, concatenated together and feed into regression module to get the rotation angle.

5.Joint Alignment Network

Joint alignment network combines the pair-wise network, a spatial transfromer and an autoencoder. Joint alignment network takes in a batch of images and feed it into pair-wise alignment network by combing each image with a reference image. Then the outputs of pair-wise alignment network are utilized to transform each image in a single batch so that aligned images are obtained. The aligned images are then feed into autoencoder to get reconstructed images. Then the reconstruction loss can be calculated.

6.Latter Procedures

After rotation angle estimated by UDL or JUDL, translation estimation can be excuted by traditional Fourier alignment such as EMAF. We have provided C++ source code to using well-trained UDL/JUDL network in EMAF.

Instructions

1.Environment for Python Source Code

The whole program is implemented in Python3 using PyTorch package. Python packages used in the program is listed in 'requirements.txt' under the source code folder. We suggest the users may use Conda to create Python environment for using the source code.

2.Binary File Usage

Under some practical situations, it may be hard to directly run the source code. Therefore, we suggest the users may use PyInstaller package to generate binary file which would be independent from Python environment. Note that CUDA version is still required to be matched with the complied environment. More instructions is provided in 'pyinstaller.txt' under the source code folder.

3.Training for New Datasets

Note that it might be hard to train the network from the very beginning for every dataset. For UDL training, we train a intialization model on the synthetic GroEL dataset and use it as initializtion for models on other datasets. Such initialization model is provided as 'initialization.pkl' under the source code folder but please note that such initialization may be not suitable for all datasets which means you may need to generate your own initialization model for training on new synthetic dataset. For JUDL training, we first train UDL as mentioned before on the dataset. Then the trained model is used as initialization for pair-wise alignment part of JUDL and continue training until convergence.

Code&Data download

Please complete the relevant information to help us improve this program and get additional support (not mandatory). Click to get the ALL_FILES. Richer comments and intructions are provided in source code to help user understand the program.

Reference: Chen, Yu-Xuan, Dagan Feng, and Hong-Bin Shen. "Cryo-EM image alignment: from pair-wise to joint with deep unsupervised difference learning." Journal of Structural Biology (2023): 107940.