Visual Storyteller

Final project for my deep learning class. Implemented a neural network capable of sequence generation conditioned on visual input.

• Input: model accepts an image
• Output: model generates a sequence of words (a sentence) describing that image.

Dataset

The local expected layout is:

data/
  Images/
  captions.txt

The image folder is large and is not committed to this repository. Extract caption_data.zip so that the repository matches the structure above before running either notebook.

Project Files

notebooks/data_and_training.ipynb  # data prep, vocabulary, train/val/test split, training
notebooks/inference.ipynb          # model loading, greedy decoding, examples, failure analysis
models/decoder_best.pt             # trained decoder weights
models/vocab.json                  # saved vocabulary
models/model_config.json           # saved model settings

Setup

Create and activate a virtual environment, then install the dependencies:

pip install -r requirements.txt

How to Run

1. Put the dataset files in data/.
2. Run notebooks/data_and_training.ipynb top to bottom to:
   • load and inspect the caption dataset
   • build the vocabulary
   • create the train/validation/test split
   • train the attention-based decoder
   • save the final model artifacts in models/
3. Run notebooks/inference.ipynb top to bottom to:
   • load the saved artifacts
   • call generate_caption(image_path, model)
   • show predictions on unseen held-out test images
   • review both good examples and failure cases

The required inference function is implemented in inference.ipynb:

def generate_caption(image_path: str, model: any) -> str:
    """Takes a path to an image and returns a generated caption string."""

Model

• Encoder: ResNet-50 pretrained on ImageNet which was frozen and the final pooling/classification layers were removed.
• Attention: Bahdanau additive attention over the encoder's spatial feature locations.
• Decoder: GRUCell-based recurrent decoder with word embeddings and greedy decoding at inference time.
• Vocabulary: 5,160 tokens with MIN_WORD_FREQ = 2.
• Split: 80% train, 10% validation, 10% test, split by image.

Training Notes

The best checkpoint was saved at epoch 5 with validation loss 2.8674. I ran 6 epochs and used the epoch 5 checkpoint because validation loss started to rise after that.

• notebooks/data_and_training.ipynb documents data loading, model definition, training progress and artifact saving.
• notebooks/inference.ipynb implements generate_caption(...), shows predictions on unseen test images and includes success/failure analysis.
• models/decoder_best.pt, models/vocab.json, and models/model_config.json are the final artifacts used by inference.