version: '2.3'
services:
  training:
    image: 172.20.40.190:5000/tf-2.0.0-gpu-py3:tfds
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=0,1
    user: "2004:2004"
    container_name: "moono_gpu_01"
    volumes:
      - "/mnt/vision-nas/moono/git-repos/stylegan2-tf-2.x:/work-dir"
      - "/mnt/vision-nas/moono/trained_models/stylegan2-tf-2.x:/model-dir"
      - "/mnt/vision-nas/data-sets/stylegan/ffhq-dataset/tfrecords/ffhq:/tfrecord-dir:ro"
    working_dir: "/work-dir"
    command: ["python", "-u", "train.py",
              "--allow_memory_growth", "false",
              "--debug_split_gpu", "false",
              "--use_tf_function", "true",
              "--use_custom_cuda", "true",
              "--model_base_dir", "/model-dir/cuda",
              "--tfrecord_dir", "/tfrecord-dir",
              "--train_res", "256",
              "--shuffle_buffer_size", "1000",
              "--batch_size_per_replica", "16"
    ]
  watching:
    image: 172.20.40.190:5000/tf-2.0.0-gpu-py3:tfds
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=1
    user: "2004:2004"
    container_name: "moono_tensorboard"
    depends_on:
      - training
    volumes:
      - "/mnt/vision-nas/moono/trained_models/stylegan2-tf-2.x:/model-dir"
    ports:
      - "36006:6006"
    working_dir: "/model-dir"
    command: ["tensorboard", "--logdir", "./", "--host", "0.0.0.0", "--port", "6006"]