Fractal Pair Similarity (FPS) Pre-training

This is an official implementation of fractal pair similarity (FPS) pre-training in On the Importance and Applicability of Pre-Training for Federated Learning.

Dependencies

• python3.6
• torch==1.10.2
• torchvision==0.11.3
• numba==0.53.1
• pytorch-lightning==1.5.10
• cv2==4.5.5

Usage

We present instructions for 1) generating FPS examples, and 2) FPS pre-training with SimSiam.

Generating FPS examples

To quicky view the images generated by FPS, please use the following command to randomly generate 50 FPS examples; each example consists of 2 fractal images.

$ python generate_fps_examples.py

The outputs will be stored in FPS_SAMPLES/. Note that image-level augmentations are also applied to the generated fractals images.

FPS pre-training with SimSiam

• For pre-training a ResNet-20 with 32x32 images.

$ bash scripts/run_fractal_simsiam_resnet20.sh ${GPUIDX} ${EXPAND} 

The ${GPUIDX} argument specifies the GPU used for training, which can usually be set to 0. The ${EXPAND} argument expands the width of the model, which is set to 1 to indicate the original ResNet-20.

• For pre-training a ResNet-18 with 224x224 images.

$ SLURM_PROCID=0 bash scripts/run_fractal_simsiam_resnet18.sh ${GPUIDX} ${PORTNUM} 

Similarly, ${GPUIDX} specifies the GPU used for training. As we use multi-gpu distributed training for ResNet-18 with 224x224 images, we need to set the ${PORTNUM} argument, which can be any unused port number (16-bit unsigned integer), like 1001. More details about how this works in pytorch can be found here.

In our experiments, we use the Slurm framework for training models on a cluster. You can modify the SLURM_PROCID environment variable and the option --ngpus-per-node to meet your own training setups if you have different needs. More details about Slurm can be found here.

Note: The above commands are for FPS pre-training with SimSiam. For training with MoCoV2, please use the scripts/run_fractal_mocov2_{resnet20|resnet18}.sh scripts instead.

Contact

If you have any questions, please contact Cheng-Hao Tu([email protected]).