docker-compose.yml

services:
  train:
    build:
      context: .
      dockerfile: docker/Dockerfile.train
    container_name: train
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      - ./results/:/code/results/
      - ./models/:/code/models/
      - ./datasets/:/code/datasets/
    command: 'python train.py --dataset fashion-mnist --model stnet --epochs 10'
    runtime: nvidia

  jupyter:
    build:
      context: .
      dockerfile: docker/Dockerfile.jupyter
    container_name: jupyter
    user: root
    ports:
      - 8888:8888
    environment:
      - JUPYTER_ENABLE_LAB=yes
      - JUPYTER_TOKEN=docker
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      - ./jupyter/:/home/jovyan/work/jupyter/
    runtime: nvidia

  train-birds:
    build:
      context: .
      dockerfile: docker/Dockerfile.train_birds
    container_name: train-birds
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      - ./results/:/code/results/
      - ./models/:/code/models/
      - ./datasets/:/code/datasets/
    command: 'python train.py --dataset birds --model stresnext --epochs 20'
    runtime: nvidia