#train.py
This program is designed for training and evaluating deep learning models to extract Gaussian parameters from absorption image data. The program supports two types of runs specified by the RUN_TYPE flag:
- Training: Trains a new model on a set of images.
- Validation: Iterates over a provided set of images and outputs metrics demonstrating model performance.
--model_name: Model name to be prepended to saved results and weights.--run_type: The type of run to execute (training,validation).--images_folder_path: Path to the folder where the atom images reside. The expected format is a.npzfile withatoms,bg, anddarkimages.--num_images: Number of images to load from the folder.--load_weight_path: Path to load stored weights for the model. Required forvalidation. (default:None)--load_results_path: Path to load the.npzfile containing previous training results. Necessary forvalidation. (default:None)--start_index: Index to begin loading images from the folder. (default:0)
--cnn_model_name: Type of convolutional neural network to use (xception,mobilenet,efficientnet,regnet). (default:mobilenet)--batch_size: Batch size for the model. (default:8)--learning_rate: Learning rate of the optimizer. (default:0.001)--save_folder_path: Path to save the weights, results, and output of the model. IfNone, defaults to./MODEL_NAME. (default:None)--zero_percent: Percentage of samples that will not generate a simulated cloud. (default:None)--num_epochs: Number of epochs to train for.--atom_sigma_detection: Standard deviations from the center of the atom cloud to the edge of the camera before deciding the simulated atoms do not exist in the image. (default:3)--num_samples: Number of Gaussian samples generated from each background provided. (default:10)--random_sample: Whether to randomly sample from the image folder path for training and validation datasets.--save_frequency: Frequency (in epochs) to save model weights and results. (default:None)--validation_frequency: Frequency (in epochs) to run validation. (default:10)--train_cloud_resample_frequency: Frequency (in epochs) to resample new Gaussians for training datasets. (default:None)--train_bg_resample_frequency: Frequency (in epochs) to resample new backgrounds for training datasets. (default:None)--train_val_folder_ratio: Ratio of total images in the folder used for training. (default:None)--train_val_ratio: Ratio of images used for training. (default:None)--n_channel: Number of images used for training/evalution. Set to 1 when using only 'Atoms' image. Set to 3 when using 'Atoms', 'Background', and 'Dark' images. (default:1)
--peak_od_ub: Upper bound of peak optical density when generating Gaussian clouds. (default:3)--peak_od_lb: Lower bound of peak optical density when generating Gaussian clouds. (default:0)--r0x_ub: Upper bound ofr0xwhen generating Gaussian clouds. (default:0.9*image.shape[1])--r0x_lb: Lower bound ofr0xwhen generating Gaussian clouds. (default:0.1*image.shape[1])--r0y_ub: Upper bound ofr0ywhen generating Gaussian clouds. (default:0.9*image.shape[0])--r0y_lb: Lower bound ofr0ywhen generating Gaussian clouds. (default:0.1*image.shape[0])--std_x_ub: Upper bound of the standard deviation of the x axis when generating Gaussian clouds (before rotation). (default:0.25*image.shape[1])--std_x_lb: Lower bound of the standard deviation of the x axis when generating Gaussian clouds (before rotation). (default:2)--std_y_ub: Upper bound of the standard deviation of the y axis when generating Gaussian clouds (before rotation). (default:0.25*image.shape[0])--std_y_lb: Lower bound of the standard deviation of the y axis when generating Gaussian clouds (before rotation). (default:2)--angle_ub: Upper bound of the CCW angle of rotation when generating Gaussian clouds. (default:pi/4)--angle_lb: Lower bound of the CCW angle of rotation when generating Gaussian clouds. (default:-pi/4)--od_offset_ub: Upper bound of the OD offset when generating Gaussian clouds. (default:0)--od_offset_lb: Lower bound of the OD offset when generating Gaussian clouds. (default:0)
- Ensure the
.npzfiles for images contain the necessary keys (atoms,bg,dark). - Adjust the
num_samples,num_epochs, and other hyperparameters according to your specific dataset and requirements.
python3 ./train.py \
--images_folder_path /path/to/images/folder \
--num_images 100 \
--model_name your_model_name \
--random_sample \
--num_epochs 70000 \
--run_type training \
--save_frequency 3000 \
--save_folder_path /path/to/save/folder \
--scale_factor 1 \
--train_cloud_resample_frequency 5 \
--train_bg_resample_frequency 20 \
--validation_frequency 1 \
--train_val_ratio 0.6 \
--train_val_folder_ratio 0.6 \
--num_samples 10 \
--batch_size 8 \
--peak_od_ub 8 \
--r0x_ub 676 \
--r0x_lb -100 \
--r0y_ub 350 \
--r0y_lb -100
python3 ./train.py \
--model_name your_model_name \
--save_folder_path ./path/to/save/folder \
--images_folder_path /path/to/images/folder \
--zero_percent 0.5 \
--num_images 10 \
--run_type validation \
--num_samples 1 \
--batch_size 1 \
--load_weight_path /path/to/model/weights.pth \
--load_results_path /path/to/model/results.npz
This program is designed for comparing CNN model against more traditional fitting technqiues. Program outputs npz file containing the results and fit times of the CNN model, 2D-LS fit, and 3x1D-LS fit.
--weights_path: The path and filename pointing to the CNN model weights (.pth file).--results_path: The path and filename pointing to the CNN model results file (.npz file).--images_path: Path to the folder where the atom images reside. The expected format is a.npzfile withatoms,bg, anddarkimages. Images may be real or simulated.
This is free software released under the terms of The MIT License. There is no warranty; not even for merchantability or fitness for a particular purpose. Consult LICENSE for copying conditions.