50 lines
2.8 KiB
Bash
50 lines
2.8 KiB
Bash
# -*- coding: utf-8 -*-
|
|
# @Author: Weisen Pan
|
|
|
|
# Initialize environment and load necessary modules
|
|
# This sets up the environment for running the necessary libraries like GCC, OpenMPI, CUDA, cuDNN, NCCL, and Python
|
|
source /etc/profile.d/modules.sh
|
|
module load gcc/11.2.0 # Load GCC version 11.2.0 for compiling
|
|
module load openmpi/4.1.3 # Load OpenMPI version 4.1.3 for distributed computing
|
|
module load cuda/11.5/11.5.2 # Load CUDA version 11.5.2 for GPU computing
|
|
module load cudnn/8.3/8.3.3 # Load cuDNN version 8.3.3 for deep learning frameworks
|
|
module load nccl/2.11/2.11.4-1 # Load NCCL version 2.11.4-1 for multi-GPU communication
|
|
module load python/3.10/3.10.4 # Load Python version 3.10.4
|
|
|
|
# Activate the Python virtual environment
|
|
# This activates the pre-configured virtual environment where necessary Python packages (e.g., PyTorch, Horovod) are installed
|
|
source ~/venv/pytorch1.11+horovod/bin/activate
|
|
|
|
# Prepare the log directory and clean up any old records
|
|
# Create a log directory for this job run and remove any previous log records
|
|
LOG_DIRECTORY="/home/projadmin/Federated_Learning/project_EdgeFLite/records/${JOB_NAME}_${JOB_ID}"
|
|
rm -rf ${LOG_DIRECTORY} # Remove old logs if they exist
|
|
mkdir -p ${LOG_DIRECTORY} # Create a new directory for current job logs
|
|
|
|
# Set up local data directory and copy dataset
|
|
# Define local data storage and copy the dataset for training the model
|
|
DATA_STORAGE="${SGE_LOCALDIR}/${JOB_ID}/"
|
|
cp -r ../summit2024/simpleFL/performance_test/cifar100/data ${DATA_STORAGE} # Copy CIFAR-100 dataset to the local directory
|
|
|
|
# Change directory to project location
|
|
# Navigate to the EdgeFLite project directory to execute the training script
|
|
cd EdgeFLite
|
|
|
|
# Execute the training process for the federated learning model
|
|
# This runs the model training with specific hyperparameters for federated learning, including architecture, dataset, and configuration settings
|
|
python train_EdgeFLite.py \
|
|
--is_fed=1 \ # Enable federated learning mode
|
|
--fixed_cluster=0 \ # Disable fixed clusters, allowing dynamic changes
|
|
--split_factor=16 \ # Set data split factor to 16
|
|
--num_clusters=6 \ # Use 6 clusters for the federated learning process
|
|
--num_selected=6 \ # Select 6 clients for each training round
|
|
--arch="wide_resnetsl16_8" \ # Use a Wide ResNet architecture with depth 16 and width 8
|
|
--dataset="cifar100" \ # Specify CIFAR-100 as the dataset
|
|
--num_classes=100 \ # CIFAR-100 has 100 output classes
|
|
--is_single_branch=0 \ # Use multi-branch (multi-head) learning
|
|
--is_amp=0 \ # Disable automatic mixed precision training
|
|
--num_rounds=650 \ # Train for 650 communication rounds
|
|
--fed_epochs=1 \ # Each client trains for 1 epoch per round
|
|
--spid="EdgeFLite_W168_96c_650r" \ # Set the unique identifier for the job
|
|
--data=${DATA_STORAGE} # Provide the location of the dataset
|