version: "3.8"
services:
  tensorflow_v1:
    build:
      context: docker
      dockerfile: tensorflow_v1.dockerfile
    image: ayo/tf_v1_benchmark
    container_name: tf_v1_benchmark
    runtime: nvidia
    volumes:
      - .:/work
    command: bash -c "umask 0000 && python benchmark.py --platform TF1 ${BENCH_ARGS}"
  tensorflow:
    build:
      context: docker
      dockerfile: tensorflow.dockerfile
    image: ayo/tf_benchmark
    container_name: tf_benchmark
    runtime: nvidia
    volumes:
      - .:/work
    command: bash -c "umask 0000 && python benchmark.py --platform TF2 TF2_V1 ${BENCH_ARGS}"
  torch:
    build:
      context: docker
      dockerfile: pytorch.dockerfile
    image: ayo/torch_benchmark
    container_name: torch_benchmark
    runtime: nvidia
    volumes:
      - .:/work
    command: bash -c "umask 0000 && python benchmark.py --platform Torch ${BENCH_ARGS}"