Kubetorch: Kubernetes에서 PyTorch ML 개발을 가속화하는 프레임워크

import kubetorch as kt
from my_repo.training_main import train

train_compute = kt.Compute(cpus="2", gpus="1")
remote_train = kt.fn(train).to(train_compute)
result = remote_train(lr=0.05, batch_size=4)

import kubetorch as kt
from my_repo.training_main import train

train_compute = kt.Compute(gpus="8").distribute("pytorch", workers=8)
remote_train = kt.fn(train).to(train_compute)
result = remote_train(lr=0.05, batch_size=4)

try:
    remote_train.train(epoch, batch_size = batch_size)
except Exception as e:
    if "CUDA out of memory" in str(e):
        batch_size = batch_size / 2
        remote_train.save()