PyTorch

Can be used with Kubeflow to orchestrate.

Kubeflow PyTorchJob

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: mnist-training
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime
              command: ["python", "/train.py"]
    Worker:
      replicas: 2
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime
              command: ["python", "/train.py"]