pai/marketplace-v2/pytorch-mnist.yaml

44 строки
1020 B
YAML

protocolVersion: 2
name: pytorch_mnist
type: job
version: 1.0
contributor: OpenPAI
description: |
This is a basic PyTorch MNIST digit recognition example running on OpenPAI.
```sh
python main.py
# CUDA_VISIBLE_DEVICES=2 python main.py # to specify GPU id to ex. 2
```
parameters:
epochs: 1
batchsize: 64
lr: 0.01
prerequisites:
- protocolVersion: 2
name: pytorch_example
type: dockerimage
version: 1.0
contributor : OpenPAI
description: |
This is an [example PyTorch Docker image on OpenPAI](https://github.com/Microsoft/pai/tree/master/examples/pytorch).
uri : openpai/pai.example.pytorch
taskRoles:
worker:
instances: 1
completion:
minSucceededInstances: 1
dockerImage: pytorch_example
resourcePerInstance:
cpu: 4
memoryMB: 8192
gpu: 1
commands:
- >
python3 examples/mnist/main.py
--epochs <% $parameters.epochs %>
--lr <% $parameters.lr %>
--batch-size <% $parameters.batchsize %>