Package: kubeml
KubeMLKubeML provides wrappers and tools that allow the interaction with user code written in PyTorch
with the distributed training and serving functionality offered by KubeMLInstallingInstall and update using pippip install kubemlUsageThe main functionality offered is in the shape of Models and Datasets. A KubeDataset is a convenience wrapper over a
torch dataset which, like when using torch, users extend with their own functionality to adapt to their data. A simple
example of how to create a dataset to train with KubeML is seen below.The Dataset classfromkubemlimportKubeDatasetfromtorchvisionimporttransformsclassMnistDataset(KubeDataset):def__init__(self):super().__init__("mnist")self.transf=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,),(0.3081,))])def__getitem__(self,index):x=self.data[index]y=self.labels[index]returnself.transf(x),y.astype('int64')def__len__(self):returnlen(self.data)The user only needs to provide to the constructor the name of the dataset as was uploaded to the KubeML storage,
the dataset will take care of fetching only the corresponding minibatches of data so that the network can be trained
with a model parallel approach.As with a normal torch dataset, the user must implement the__getitem__and__len__methods to iterate over the dataset.
The dataset exposes two member variables:dataHolds the features used as input to the networklabelsHolds the output labelsBoth are saved as numpy arrays.The Model classThe other main component is the model class. This abstract class abstracts the complexity of distributing the training
among multiple workers, nodes and GPUs. The constructor only takes a torch model and the dataset as a parameter. The user only needs
to implement the abstract methods of the class,train,infer,validateinitandconfigure_optimizerswith the behavior they
want from the network.The Kubenet exposes thebatch_sizeandlrarguments which the user can change when starting the train jobfromkubemlimportKubeModelimporttorchimporttorch.nnasnnimportnumpyasnpclassKubeLeNet(KubeModel):def__init__(self,network,dataset):super().__init__(network,dataset,gpu=True)@abstractmethoddefconfigure_optimizers(self)->torch.optim.Optimizer:pass# Train trains the model for an epoch and returns the loss@abstractmethoddeftrain(self,x,y,batch_index)->float:pass# Validate validates the model on the test data and returns a tuple# of (accuracy, loss)@abstractmethoddefvalidate(self,x,y,batch_index)->Tuple[float,float]:pass# Infer receives the data points or images as a list and returns# the predictions of the network@abstractmethoddefinfer(self,data:List[Any])->Union[torch.Tensor,np.ndarray,List[float]]:pass# Init initializes the model in a particular way@abstractmethoddefinit(self,model:nn.Module):passAn example implementation of theinitandtrainfunctions can be done as follows# Train trains the model for an epoch and returns the lossdeftrain(self,x,y,batch_index)->float:# define the device for training and load the dataloss_fn=nn.CrossEntropyLoss()total_loss=0self.optimizer.zero_grad()output=self(x)# compute loss and backprop# logging.debug(f'Shape of the output is {output.shape}, y is {y.shape}')loss=loss_fn(output,y)loss.backward()# step with the optimizerself.optimizer.step()total_loss+=loss.item()ifbatch_index%10==0:logging.info(f"Index{batch_index}, error:{loss.item()}")returntotal_loss# Intialize the network as a pytorch modeldefinit(self,model):definit_weights(m:nn.Module):ifisinstance(m,nn.Conv2d):nn.init.xavier_uniform_(m.weight)nn.init.constant_(m.bias,0.01)ifisinstance(m,nn.Linear):nn.init.xavier_uniform_(m.weight)nn.init.constant_(m.bias,0.01)model.apply(init_weights)Writing the training functionAt the moment of creating a serverless function which will serve as a worker for the model training process, the
steps are simple, simply write the code initializing the network in themainmethod of the function, and callstarton the KubeML model.defmain():# Create the PyTorch Modellenet=LeNet()dataset=MnistDataset()kubenet=KubeLeNet(lenet,dataset)returnkubenet.start()