How santiagosabogall/gan-for-galaxy-image-generation works
1. This repository is a single-purpose TensorFlow project for training a Wasserstein GAN with Gradient Penalty (WGAN-GP) to generate 64×64 RGB galaxy images from the Galaxy Zoo 2 dataset. It is an offline model-training pipeline, not a serving or application backend.
2. The code is split into five clear subsystems under src/:
- config.py for global hyperparameters and paths
- data.py for dataset acquisition and preprocessing
- models.py for Generator and Critic definitions
- train.py for the custom WGAN-GP training loop
- utils.py for visualization, GIF creation, and plotting
3. The main execution entrypoint is python -m src.train. That script orchestrates the full lifecycle: load configuration, obtain data, instantiate models, run alternating Critic/Generator updates, and emit artifacts during training.
4. The data boundary is external: the project downloads the Galaxy Zoo 2 dataset via kagglehub and feeds it through a tf.data pipeline. This means the training loop consumes an abstracted TensorFlow dataset rather than reading raw files directly.
5. The preprocessing boundary likely normalizes images into the range expected by the Generator/Critic pair. The README indicates a tanh output on the Generator, so the image pipeline almost certainly maps real images into a matching normalized scale for WGAN training.
6. The model subsystem contains two networks with distinct roles:
- Generator: latent vector z ∈ R^128 → synthetic 64×64×3 image
- Critic: 64×64×3 image → scalar Wasserstein score
The architecture uses transposed convolutions in the Generator and strided convolutions in the Critic.
7. The WGAN-GP-specific architectural constraint is important: the Critic intentionally avoids BatchNormalization, because gradient penalty depends on well-behaved sample-wise gradients. That choice is a key boundary between “ordinary CNN classifier” design and WGAN critic design.
8. The training subsystem is a custom adversarial optimization loop rather than a standard Keras fit() workflow. It alternates Critic and Generator updates and computes the gradient penalty term explicitly, which is central to the repository’s behavior.
9. The repository also has a clear artifact/output boundary. During training, it writes:
- model checkpoints into checkpoints/
- preview images into output/
- TensorBoard logs into logs/
These outputs are the main integration points for monitoring and resuming experiments.
10. utils.py provides the post-training/monitoring layer: it turns epoch snapshots into visual progress assets, including an animated GIF that summarizes training evolution. This is a presentation/inspection concern, separate from the core optimization logic.
11. requirements.txt and the README indicate the runtime stack is Python 3.9+ with TensorFlow 2.15+. External tooling that materially affects the architecture includes TensorBoard for metrics inspection and kagglehub for dataset retrieval/caching.
12. The notebook GAN_images_of_galaxies.ipynb is a legacy or exploratory artifact, while src/ is the productionized training implementation. In an architecture diagram, the notebook should be shown as optional historical context, not part of the main runtime path.
13. In diagram terms, the primary flow is:
Kaggle/Galaxy Zoo 2 → data pipeline → training loop → Generator/Critic updates → checkpoints/logs/images/GIFs.
The only persistent external dependencies are the dataset source and the local artifact directories used for experiment tracking.
Main entry points: train.py WGAN-GP loop · data.py Data pipeline · models.py Models · utils.py Utils
Components 18
Source code
Core TensorFlow training implementation under src/
- Train entry · Main orchestration entrypoint for WGAN-GP training
src/train.py - Config · Global hyperparameters, paths, and run settings
src/config.py - Data pipeline · Dataset acquisition, preprocessing, and tf.data construction
src/data.py - Models · Generator and critic definitions
src/models.py - Utils · Plotting, previews, and GIF creation
src/utils.py - Generator · Maps latent vectors to 64x64 RGB galaxy images
src/models.py - Critic · Scores real and generated images for Wasserstein training
src/models.py - WGAN-GP loop · Alternating critic and generator updates with custom training logic
src/train.py - Gradient penalty · Explicit WGAN-GP regularization computed in the loop
src/train.py - Notebook · Exploratory or historical implementation artifact
GAN_images_of_galaxies.ipynb
External dependencies
Dataset retrieval and experiment tooling outside the repo
- KaggleHub · Downloads and caches the Galaxy Zoo 2 dataset
- tf.data · TensorFlow dataset abstraction consumed by training
- Galaxy Zoo 2 · Source dataset of galaxy images
- TensorBoard · Inspects training summaries and logged scalars
Training artifacts
Files and directories produced during training
- Checkpoints · Saved model state for resuming training
- Output images · Preview images generated during training
- TensorBoard logs · Scalar and summary data for monitoring
- Training GIF · Animated summary of training progress
Connections 22
- Train entry → Config: loads · Reads hyperparameters, paths, and training settings
- Train entry → Data pipeline: prepares · Builds the dataset input pipeline
- Data pipeline → KaggleHub: downloads · Fetches Galaxy Zoo 2 through KaggleHub
- KaggleHub → Galaxy Zoo 2: sources · External dataset source
- Data pipeline → tf.data: yields · Produces TensorFlow datasets for training
- Train entry → Models: instantiates · Creates generator and critic networks
- Models → Generator: defines · Generator architecture and output scaling
- Models → Critic: defines · Critic architecture without batch normalization
- Train entry → WGAN-GP loop: runs · Executes the custom adversarial optimization loop
- WGAN-GP loop → tf.data: consumes · Reads batches from the dataset pipeline
- WGAN-GP loop → Critic: updates · Trains critic on real and generated images
- WGAN-GP loop → Generator: updates · Trains generator to fool the critic
- WGAN-GP loop → Gradient penalty: computes · Adds WGAN-GP regularization to the critic loss
- WGAN-GP loop → Checkpoints: writes · Persists model checkpoints
- WGAN-GP loop → TensorBoard logs: writes · Emits TensorBoard summaries
- WGAN-GP loop → Output images: writes · Saves preview images during training
- TensorBoard logs → TensorBoard: viewed in · Monitored with TensorBoard
- Utils → Output images: renders · Creates visual samples and progress images
- Utils → Training GIF: assembles · Combines epoch snapshots into an animation
- Train entry → Utils: uses · Calls visualization and artifact helpers
- Notebook → Models: prototypes · Legacy exploratory model code path
- Notebook → Data pipeline: prototypes · Legacy exploratory data handling