Image Patch SSL (TensorFlow)

Self-supervised patch-level model with a shared latent z, reconstruction decoder, and contrastive objective with EMA target encoder/projector.

Project layout

• src/data.py
• src/models.py
• src/train.py
• src/eval.py
• checkpoints/
• logs/

Setup

1. Python 3.9–3.11 recommended. GPU optional but recommended.
2. Install deps:
   • PowerShell:

     python -m pip install -r requirements.txt

Data

• Place images under data/ or any folder. Supported by TensorFlow decode (jpg/png/etc.).
• Images are resized to 1024×1024 on-the-fly.
• From each image, num_patches_per_image random 16×16 patches are sampled.

Train

Run from the project root using the module form (so imports resolve):

python -m src.train --file_pattern "data/*" --epochs 2 --batch_size 64 --cache

Examples (Windows globbing):

• All jpg recursively: --file_pattern "data/**/*.jpg"
• Single folder: --file_pattern "data/*"

Key flags:

• --patch_size 16 (default)
• --resize 1024
• --latent_dim 256
• --proj_dim 128
• --pred_hidden 128
• --tau 0.2
• --lambda_contrast 0.2
• --momentum 0.996
• --batch_size 256 (reduce if OOM)
• --mixed_precision to enable float16 policy on supported GPUs

TensorBoard:

tensorboard --logdir logs

Evaluate

python -m src.eval --file_pattern "data/*" --max_steps 50

Prints MSE/PSNR/SSIM and a simple Retrieval@1 on two augmented views.

Notes

• For large datasets, consider creating TFRecords of resized images.
• EMA momentum can be annealed upward during training.
• If recon loss is high, consider milder augs or larger latent_dim.