Attention-OCR Training Template

attention-ocr-training.yaml

arguments:
  parameters:
  - name: source
    value: https://github.com/tensorflow/models.git
    displayName: Model source code
    type: hidden
    visibility: private

  - name: trainingsource
    value: https://github.com/onepanelio/LicensePlateOcr.git
    type: hidden
    visibility: private
    
  - name: batch-size
    value: 1
    visibility: internal
    
  - name: max-steps
    value: 11000
    visibility: internal
    
  - name: cvat-annotation-path
    value: annotation-dump/license-plate-shorter-new-15
    displayName: Dataset path
    hint: Path to annotated data in default object storage (i.e S3). In CVAT, this parameter will be pre-populated.
    visibility: internal

  - name: cvat-output-path
    value: workflow-data/output/license-plate-ocr-output1
    hint: Path to store output artifacts in default object storage (i.e s3). In CVAT, this parameter will be pre-populated.
    displayName: Workflow output path
    visibility: internal
    

  - name: tf-image
    value: tensorflow/tensorflow:1.13.1-py3
    type: select.select
    displayName: Select TensorFlow image
    visibility: public
    hint: Select the GPU image if you are running on a GPU node pool
    options:
    - name: 'TensorFlow 1.13.1 CPU Image'
      value: 'tensorflow/tensorflow:1.13.1-py3'
    - name: 'TensorFlow 1.13.1 GPU Image'
      value: 'tensorflow/tensorflow:1.13.1-gpu-py3'
    

  - displayName: Node pool
    hint: Name of node pool or group to run this workflow task
    type: select.select
    name: sys-node-pool
    value: Standard_D4s_v3
    visibility: public
    required: true
    options:
    - name: 'CPU: 2, RAM: 8GB'
      value: Standard_D2s_v3
    - name: 'CPU: 4, RAM: 16GB'
      value: Standard_D4s_v3
    - name: 'GPU: 1xK80, CPU: 6, RAM: 56GB'
      value: Standard_NC6
    - name: 'GPU: 1xV100'
      value: Standard_NC6s_v3
      
  - name: dump-format
    value: cvat_tfrecord
    visibility: public
    
entrypoint: main
templates:
- dag:
    tasks:
    - name: train-model
      template: tensorflow

  name: main
- container:
    args:
    - |
      apt-get update && \
      apt-get install -y python3-pip git wget unzip libglib2.0-0 libsm6 libxext6 libxrender-dev && \
      
      cd /mnt/src/tf/research && \
      export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim && \
      cd /mnt/src/train && \
      pip install -r requirements.txt && \
      cp -f custom.py /mnt/src/tf/research/attention_ocr/python/datasets/ && \
      cp -f __init__.py /mnt/src/tf/research/attention_ocr/python/datasets/ && \
      cp -f ./data/charset_size.txt /mnt/data/datasets/ && \
      python generate_tfrecords.py \
        --charset_path=/mnt/src/train/data/charset_size.txt \
        --data_dir=/mnt/data/datasets/ \
        --output_tfrecord=/mnt/data/datasets/tfexample_train  \
        --text_length=14 && \
      
      cd /mnt/src/tf/research/attention_ocr/python/ && \
      wget http://download.tensorflow.org/models/attention_ocr_2017_08_09.tar.gz && \
      tar xf attention_ocr_2017_08_09.tar.gz && \
      export PYTHONPATH=$PYTHONPATH:./datasets/ && \
      python train.py \
        --dataset_name=custom \
        --batch_size={{workflow.parameters.batch-size}} \
        --train_log_dir=/mnt/output/ \
        --checkpoint=model.ckpt-399731 \
        --max_number_of_steps={{workflow.parameters.max-steps}}
    command:
    - sh
    - -c
    image: '{{workflow.parameters.tf-image}}'
    volumeMounts:
    - mountPath: /mnt/data
      name: data
    - mountPath: /mnt/output
      name: output
    workingDir: /mnt/src
  nodeSelector:
    beta.kubernetes.io/instance-type: '{{workflow.parameters.sys-node-pool}}'
  inputs:
    artifacts:
    - name: data
      path: /mnt/data/datasets/
      s3:
        key: '{{workflow.namespace}}/{{workflow.parameters.cvat-annotation-path}}'
   
    - git:
        repo: '{{workflow.parameters.source}}'
      name: src
      path: /mnt/src/tf
    - git:
        repo: '{{workflow.parameters.trainingsource}}'
      name: tsrc
      path: /mnt/src/train
  name: tensorflow
  outputs:
    artifacts:
    - name: model
      optional: true
      path: /mnt/output
      s3:
        key: '{{workflow.namespace}}/{{workflow.parameters.cvat-output-path}}/{{workflow.name}}'

volumeClaimTemplates:
- metadata:
    creationTimestamp: null
    name: data
  spec:
    accessModes:
    - ReadWriteOnce
    resources:
      requests:
        storage: 200Gi
- metadata:
    creationTimestamp: null
    name: output
  spec:
    accessModes:
    - ReadWriteOnce
    resources:
      requests:
        storage: 200Gi