Unsupervised pre-training on a large dataset of cellular EM images

It’s well-known in deep learning (DL) that pre-training on a very large dataset followed by transfer learning to a smaller dataset is extremely effective. In fact, it’s a cornerstone of many state-of-the-art DL algorithms which use weights pre-trained on ImageNet for other classification, object detection, and segmentation tasks. If you followed deep learning research last year, you’ll have noticed an explosion in new algorithms for unsupervised pre-training (MoCo from Facebook, SimCLR from Google, BYOL from DeepMind). With these advances, it’s possible to pre-train CNNs without ANY upfront annotations and those pre-trained weights will be better for transfer learning.

We used the MoCo algorithm and a newly constructed dataset that we call CEM500K to test whether models pre-trained on EM images would be better for EM segmentation tasks than models pre-trained on ImageNet. Unsurprisingly, they were. Using a few publicly available benchmarks we showed that CEM500K-moco pre-training was substantially better than random initialization and pre-training on ImageNet (IN). Moreover, the pre-trained models (which are vanilla 2D U-Nets with a ResNet50 encoder) outperformed previous state-of-the-art algorithms for all benchmarks on which we could make a fair comparison.

Our preprint, code, CEM500K dataset and pre-trained weights are all publicly available. There are more details and links to everything at GitHub - volume-em/cellemnet: Utilities for using CellEMNet for unsupervised pretraining and downstream benchmark evaluation protocols..

Even if EM isn’t your domain, we hope that this draws your attention to the work being done with unsupervised pre-training. Just this week, Google Research posted a paper that uses the SimCLR algorithm to set new state-of-the-art results for medical image classification [2101.05224v1] Big Self-Supervised Models Advance Medical Image Classification.

9 Likes

This looks really useful! Do you think the weights would be better than ImageNet for other imaging methods (e.g. confocal microscopy)?

Thank you! I would be surprised if they were much better for confocal microscopy than ImageNet but it’s difficult to know exactly what a model learns during pre-training. Still, building a new dataset that’s domain specific for another imaging method is much easier when you don’t need any labels. All you need is to gather as many images as you can find, drop them in a folder and run the MoCo training script. That script is in the Github repo along with all the scripts for evaluating the pre-trained weights on a segmentation task.