Merge pull request #12 from berenslab/feature-more-image-transforms

Feature more image transforms

.github/workflows/config_scan.yml

@@ -29,8 +29,17 @@ jobs:
         singularity registry login --username ${{ github.actor }} --password ${{ secrets.GITHUB_TOKEN }} oras://ghcr.io
         singularity pull ${{ env.sif_file }} ${{ env.singularity_image }}
 
-    - name: Scan classification config
+    - name: Prepare experiment files
       run: |
         cp -r resources/config_templates/* config/
-        singularity exec ${{ env.sif_file }} \
-        python main.py -m +experiment=cifar10-class-recon command=scan system.device=cpu
+        for file in config/user/experiment/*.yaml; do
+          filename=$(basename "$file" .yaml)
+          echo "$filename" >> experiments.txt
+        done
+
+    - name: Scan classification configs
+      run: |
+        while IFS= read -r experiment; do
+          singularity exec ${{ env.sif_file }} \
+          python main.py +experiment="$experiment" command=scan system.device=cpu
+        done < experiments.txt

.github/workflows/container_build.yml

@@ -29,9 +29,9 @@ jobs:
     - name: Scan classification config / ensure minimal functionality
       run: |
         cp -r resources/config_templates/* config/
-        singularity exec retinal-rl.sif python main.py -m +experiment=cifar10-class-recon command=scan system.device=cpu
+        singularity exec retinal-rl.sif python main.py -m +experiment=classification command=scan system.device=cpu
 
     - name: Push to ghcr.io
       run: |
         singularity registry login --username ${{ github.actor }} --password ${{ secrets.GITHUB_TOKEN }} oras://ghcr.io
-        singularity push retinal-rl.sif oras://ghcr.io/berenslab/retinal-rl:singularity-image-latest
+        singularity push retinal-rl.sif oras://ghcr.io/berenslab/retinal-rl:singularity-image-latest

main.py

@@ -21,6 +21,7 @@
 from runner.train import train
 from runner.util import delete_results
 
+# Load the eval resolver for OmegaConf
 OmegaConf.register_new_resolver("eval", eval)
 
 

resources/config_templates/user/brain/deep-autoencoder.yaml

@@ -0,0 +1,92 @@
+# Configuration for a retinal classifier neural network
+
+# Define the input sensor
+sensors:
+  vision: # Vision is a list of 3 elements
+    - 3  # Number of color channels
+    - 160  # Height of the input image
+    - 160  # Width of the input image
+
+# Define how the various NeuralCircuits are connected to each other
+connections:
+  - ["vision", "retina"]  # Input to retina
+  - ["retina","thalamus"]  # Retina to thalamus
+  - ["thalamus","visual_cortex"]  # Thalamus to visual cortex
+  - ["visual_cortex", "inferotemporal"]  # Visual cortex to inferotemporal cortex
+  - ["inferotemporal", "prefrontal"] 
+  - ["prefrontal", "classifier"]  # Prefrontal cortex to classifier
+  - ["inferotemporal", "inferotemporal_decoder"]  # prefrontal to FC decoder
+  - ["inferotemporal_decoder", "decoder"]  # FC decoder to convoluational decoder
+
+# Define the individual nodes (neural circuits) of the network. Many circuit
+# parameters are interpolated from the experiment config.
+circuits:
+  # Retina: initial processing of visual input
+  retina:
+    _target_: retinal_rl.models.circuits.convolutional.ConvolutionalEncoder
+    num_layers: 2
+    num_channels: [16,32]  # Two layers with 16 and 32 channels
+    kernel_size: 8
+    stride: 2
+    act_name: ${activation}
+    layer_names: ["bipolar", "retinal_ganglion"]  # Names inspired by retinal cell types
+
+  # Thalamus: relay and processing station
+  thalamus:
+    _target_: retinal_rl.models.circuits.convolutional.ConvolutionalEncoder
+    num_layers: 1
+    num_channels: 64
+    kernel_size: 5
+    stride: 1
+    act_name: ${activation}
+    layer_names: ["lgn"]  # Lateral Geniculate Nucleus
+
+  # Visual Cortex: higher-level visual processing
+  visual_cortex:
+    _target_: retinal_rl.models.circuits.convolutional.ConvolutionalEncoder
+    num_layers: 1
+    num_channels: 64
+    kernel_size: 8
+    stride: 2
+    act_name: ${activation}
+    layer_names: ["v1"]  # Primary Visual Cortex
+
+  # Inferotemporal Cortex: Associations
+  inferotemporal:
+    _target_: retinal_rl.models.circuits.fully_connected.FullyConnectedEncoder
+    output_shape:
+      - 64  # Size of the latent representation
+    hidden_units: 
+      - 64  # Number of hidden units
+    act_name: ${activation}
+
+  # Prefrontal Cortex: high-level cognitive processing
+  prefrontal:
+    _target_: retinal_rl.models.circuits.fully_connected.FullyConnectedEncoder
+    output_shape:
+      - 32  # Size of the latent representation
+    hidden_units: 
+      - 32  # Number of hidden units
+    act_name: ${activation}
+
+  # Prefrontal Cortex: high-level cognitive processing
+  inferotemporal_decoder:
+    _target_: retinal_rl.models.circuits.fully_connected.FullyConnectedDecoder
+    output_shape: "inferotemporal.input_shape"  # Size of the latent representation
+    hidden_units: 
+      - 64  # Number of hidden units
+    act_name: ${activation}
+
+  # Decoder: for reconstructing the input from the latent representation
+  decoder:
+    _target_: retinal_rl.models.circuits.convolutional.ConvolutionalDecoder
+    num_layers: 4
+    num_channels: [64,32,16,3]  # For a symmetric encoder, this should be the reverse of the num_channels in the CNN layers up to the point of decoding (in this case, the thalamus)
+    kernel_size: [5,8,8,8]
+    stride: [1,2,2,2]
+    act_name: ${activation}
+
+  # Classifier: for categorizing the input into classes
+  classifier:
+    _target_: retinal_rl.models.circuits.task_head.linear_classifier.LinearClassifier
+    num_classes: 10  # Number of output classes

resources/config_templates/user/brain/shallow-autoencoder.yaml

@@ -12,9 +12,9 @@ connections:
   - ["vision", "retina"]  # Input to retina
   - ["retina","thalamus"]  # Retina to thalamus
   - ["thalamus","visual_cortex"]  # Thalamus to visual cortex
-  - ["visual_cortex", "pfc"]  # Visual cortex to prefrontal cortex
-  - ["pfc", "classifier"]  # Prefrontal cortex to classifier
-  - ["thalamus", "decoder"]  # Thalamus to decoder (for reconstruction)
+  - ["visual_cortex", "prefrontal"]  # Visual cortex to prefrontal cortex
+  - ["prefrontal", "classifier"]  # Prefrontal cortex to classifier
+  - ["visual_cortex", "decoder"]  # Thalamus to decoder (for reconstruction)
 
 # Define the individual nodes (neural circuits) of the network. Many circuit
 # parameters are interpolated from the experiment config.
@@ -24,7 +24,9 @@ circuits:
     _target_: retinal_rl.models.circuits.convolutional.ConvolutionalEncoder
     num_layers: 2
     num_channels: [16,32]  # Two layers with 16 and 32 channels
-    kernel_size: 8
+    kernel_size: 
+      - ${bp_kernel_size}
+      - ${rgc_kernel_size}
     stride: 2
     act_name: ${activation}
     layer_names: ["bipolar", "retinal_ganglion"]  # Names inspired by retinal cell types
@@ -34,7 +36,7 @@ circuits:
     _target_: retinal_rl.models.circuits.convolutional.ConvolutionalEncoder
     num_layers: 1
     num_channels: 64
-    kernel_size: 5
+    kernel_size: ${lgn_kernel_size}
     stride: 1
     act_name: ${activation}
     layer_names: ["lgn"]  # Lateral Geniculate Nucleus
@@ -44,13 +46,13 @@ circuits:
     _target_: retinal_rl.models.circuits.convolutional.ConvolutionalEncoder
     num_layers: 1
     num_channels: 64
-    kernel_size: 8
+    kernel_size: ${v1_kernel_size}
     stride: 2
     act_name: ${activation}
     layer_names: ["v1"]  # Primary Visual Cortex
 
   # Prefrontal Cortex: high-level cognitive processing
-  pfc:
+  prefrontal:
     _target_: retinal_rl.models.circuits.fully_connected.FullyConnectedEncoder
     output_shape:
       - 128  # Size of the latent representation
@@ -61,10 +63,14 @@ circuits:
   # Decoder: for reconstructing the input from the latent representation
   decoder:
     _target_: retinal_rl.models.circuits.convolutional.ConvolutionalDecoder
-    num_layers: 3
-    num_channels: [32,16,3]  # For a symmetric encoder, this should be the reverse of the num_channels in the CNN layers up to the point of decoding (in this case, the thalamus)
-    kernel_size: [5,8,8]
-    stride: [1,2,2]
+    num_layers: 4
+    num_channels: [64,32,16,3]  # For a symmetric encoder, this should be the reverse of the num_channels in the CNN layers up to the point of decoding (in this case, the thalamus)
+    kernel_size: 
+      - ${v1_kernel_size}
+      - ${lgn_kernel_size}
+      - ${rgc_kernel_size}
+      - ${bp_kernel_size}
+    stride: [2,1,2,2]
     act_name: ${activation}
 
   # Classifier: for categorizing the input into classes

resources/config_templates/user/dataset/cifar10.yaml

@@ -0,0 +1,31 @@
+# Configuration for the CIFAR10 dataset in classification experiments.
+
+name: cifar10  # Name of the dataset
+
+imageset:
+  _target_: retinal_rl.dataset.Imageset
+  source_transforms:
+    - _target_: retinal_rl.datasets.transforms.ScaleShiftTransform
+      vision_width: 160
+      vision_height: 160
+      image_rescale_range: [1, 5]
+  noise_transforms:
+    - _target_: retinal_rl.datasets.transforms.ShotNoiseTransform
+      lambda_range:
+        - ${eval:"0.5 if ${shot_noise_transform} else 0"}
+        - ${eval:"1.5 if ${shot_noise_transform} else 0"}
+    - _target_: retinal_rl.datasets.transforms.ContrastTransform
+      contrast_range:
+        - ${eval:"0.01 if ${contrast_noise_transform} else 1"}
+        - ${eval:"1.2 if ${contrast_noise_transform} else 1"}
+    - _target_: retinal_rl.datasets.transforms.IlluminationTransform
+      brightness_range:
+        - ${eval:"0.1 if ${brightness_noise_transform} else 1"}
+        - ${eval:"1.5 if ${brightness_noise_transform} else 1"}
+    - _target_: retinal_rl.datasets.transforms.BlurTransform
+      blur_range:
+        - ${eval:"0 if ${blur_noise_transform} else 0"}
+        - ${eval:"2 if ${blur_noise_transform} else 0"}
+  apply_normalization: true
+  fixed_transformation: false
+  multiplier: 1

resources/config_templates/user/experiment/cifar10-class-recon.yaml → resources/config_templates/user/experiment/classification.yaml

@@ -1,9 +1,9 @@
 # @package _global_
 defaults:
   - _self_
-  - override /dataset: cifar10-decontrast
+  - override /dataset: cifar10
   - override /brain: shallow-autoencoder
-  - override /optimizer: recon-weight
+  - override /optimizer: class-recon
 
 # This is the main entry point for control of a retinal-rl experiment. Variables
 # created here will be top-level, and defaults can be set for the various parts
@@ -15,8 +15,17 @@ framework: classification
 # in sweep, dataset, brain, and optimizer. A major use for this is interpolating
 # values in the subconfigs, and then looping over them in a sweep.
 activation: "elu"
-activation_sparsity: 0.0001
-weight_decay: 0.0001
-sparse_objective: retinal_rl.models.objective.L1Sparsity
+
+bp_kernel_size: 14
+rgc_kernel_size: 10
+lgn_kernel_size: 5
+v1_kernel_size: 8
+
 recon_weight_retina: 1
-recon_weight_thalamus: 0.99
+recon_weight_thalamus: 0.99
+recon_weight_cortex: 0.9
+
+shot_noise_transform: True
+contrast_noise_transform: True
+brightness_noise_transform: True
+blur_noise_transform: True

resources/config_templates/user/optimizer/recon-weight.yaml → resources/config_templates/user/optimizer/class-recon.yaml

@@ -21,6 +21,7 @@ goal:
         weight: 1
     target_circuits: # Circuit parameters to optimize with this optimizer.  We train the retina and the decoder exclusively to maximize reconstruction
       - decoder
+      - inferotemporal_decoder
   mixed:
     min_epoch: 0
     max_epoch: 100
@@ -31,6 +32,17 @@ goal:
         weight: ${eval:'1-${recon_weight_thalamus}'}
     target_circuits: # The thalamus is somewhat sensitive to task losses
       - thalamus
+  cortex:
+    min_epoch: 0
+    max_epoch: 100
+    losses:
+      - _target_: retinal_rl.models.loss.ReconstructionLoss
+        weight: ${recon_weight_cortex}
+      - _target_: retinal_rl.classification.loss.ClassificationLoss
+        weight: ${eval:'1-${recon_weight_cortex}'}
+    target_circuits: # Visual cortex and downstream layers are driven by the task
+      - visual_cortex
+      - inferotemporal
   class:
     min_epoch: 0
     max_epoch: 100
@@ -40,6 +52,5 @@ goal:
       - _target_: retinal_rl.classification.loss.PercentCorrect
         weight: 0
     target_circuits: # Visual cortex and downstream layers are driven by the task
-      - visual_cortex
-      - pfc
+      - prefrontal
       - classifier

resources/config_templates/user/sweep/kernel-size.yaml → resources/config_templates/user/sweep/recon-weight-sweep.yaml

@@ -11,5 +11,9 @@ method: grid
 project: retinal-rl
 
 parameters:
-  kernel_size:
-    values: [6,8,10]
+  recon_weight_retina: 
+    values: [0,0.9,0.99,0.999,0.9999,1]
+  recon_weight_thalamus: 
+    values: [0,0.9,0.99,0.999,0.9999,1]
+  optimizer:
+    value: "recon-weight"

resources/config_templates/user/sweep/shallow-transform-sweep.yaml

@@ -0,0 +1,27 @@
+# This is a standard template for a wandb sweep.
+# See: https://docs.wandb.ai/guides/sweeps/define-sweep-configuration
+#
+# Sweeps can be initialized with retinal-rl from the command line with:
+#
+# python +experiment=experiment_name command=sweep
+
+description: Evaluating the effect of different combinations of transform.
+program: main.py
+method: grid
+project: retinal-rl
+
+parameters:
+  activation:
+    value: "elu"
+  optimizer:
+    value: "recon-weight"
+  brain:
+    value: "shallow-autoencoder"
+  shot_noise_transform:
+    values: [False,True]
+  contrast_noise_transform:
+    values: [False,True]
+  brightness_noise_transform:
+    values: [False,True]
+  blur_noise_transform:
+    values: [False,True]