Package: mlfab
mlfabWhat is this?This is a framework for trying out machine learning ideas.Getting StartedInstall the package using:pipinstallmlfabCreate a new project quickly using the command line tool:create-mlfab-projectname(--descriptionDESCRIPTION)(--authorAUTHOR)(--emailEMAIL)(--urlURL)Simple ExampleThis framework provides an abstraction for quickly implementing and training PyTorch models. The workhorse for doing this ismlfab.Task, which wraps all of the training logic into a single cohesive unit. We can override functions on that method to get special functionality, but the default functionality is often good enough. Here's an example for training an MNIST model:fromdataclassesimportdataclassimporttorch.nn.functionalasFfromdpshdl.datasetimportDatasetfromdpshdl.impl.mnistimportMNISTfromtorchimportTensor,nnfromtorch.optim.optimizerimportOptimizerimportmlfab@dataclassclassConfig(mlfab.Config):in_dim:int=mlfab.field(1,help="Number of input dimensions")classMnistClassification(mlfab.Task[Config]):def__init__(self,config:Config)->None:super().__init__(config)self.model=nn.Sequential(nn.Conv2d(config.in_dim,32,3,padding=1),nn.BatchNorm2d(32),nn.ReLU(),nn.Conv2d(32,32,3,padding=1),nn.BatchNorm2d(32),nn.ReLU(),nn.MaxPool2d(2),nn.Conv2d(32,64,3,padding=1),nn.BatchNorm2d(64),nn.ReLU(),nn.Conv2d(64,64,3,padding=1),nn.BatchNorm2d(64),nn.ReLU(),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(64*7*7,128),nn.BatchNorm1d(128),nn.ReLU(),nn.Linear(128,10),)defset_loggers(self)->None:self.add_logger(mlfab.StdoutLogger(),mlfab.TensorboardLogger(self.exp_dir),)defget_dataset(self,phase:mlfab.Phase)->Dataset[tuple[Tensor,Tensor]]:root_dir=mlfab.get_data_dir()/"mnist"returnMNIST(root_dir=root_dir,train=phase=="train")defbuild_optimizer(self)->Optimizer:returnmlfab.Adam.get(self,lr=1e-3)defforward(self,x:Tensor)->Tensor:returnself.model(x)defget_loss(self,batch:tuple[Tensor,Tensor],state:mlfab.State)->Tensor:x,y=batchyhat=self(x)self.log_step(batch,yhat,state)returnF.cross_entropy(yhat,y.squeeze(-1))deflog_valid_step(self,batch:tuple[Tensor,Tensor],output:Tensor,state:mlfab.State)->None:(x,y),yhat=batch,outputdefget_label_strings()->list[str]:ytrue,ypred=y.squeeze(-1),yhat.argmax(-1)return[f"ytrue={ytrue[i]}, ypred={ypred[i]}"foriinrange(len(ytrue))]self.log_labeled_images("images",lambda:(x,get_label_strings()))if__name__=="__main__":# python -m examples.mnistMnistClassification.launch(Config(batch_size=16))Let's break down each part individually.ConfigTasks are parametrized using a config dataclass. Theml.fieldfunction is a lightweight wrapper arounddataclasses.fieldwhich is a bit more ergonomic, andml.Configis a bigger dataclass which contains a bunch of other options for configuring training.@dataclassclassConfig(mlfab.Config):in_dim:int=mlfab.field(1,help="Number of input dimensions")ModelAll tasks should subclassml.Taskand override the genericConfigwith the task-specific config. This is very important, not just because it makes your life easier by working nicely with your typechecker, but because the framework looks at the generic type when resolving the config for the given task.classMnistClassification(mlfab.Task[Config]):def__init__(self,config:Config)->None:super().__init__(config)self.model=nn.Sequential(nn.Conv2d(config.in_dim,32,3,padding=1),nn.BatchNorm2d(32),nn.ReLU(),nn.Conv2d(32,32,3,padding=1),nn.BatchNorm2d(32),nn.ReLU(),nn.MaxPool2d(2),nn.Conv2d(32,64,3,padding=1),nn.BatchNorm2d(64),nn.ReLU(),nn.Conv2d(64,64,3,padding=1),nn.BatchNorm2d(64),nn.ReLU(),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(64*7*7,128),nn.BatchNorm1d(128),nn.ReLU(),nn.Linear(128,10),)Loggersmlfabsupports logging to multiple downstream loggers, and provides a bunch of helper functions for doing common logging operations, like rate limiting, converting image resolution to normal sizes, overlaying captions on images, and more.If this function is not overridden, the task will just log tostdout.defset_loggers(self)->None:self.add_logger(mlfab.StdoutLogger(),mlfab.TensorboardLogger(self.exp_dir),)DatasetsThe task should return the dataset used for training, based on the phase.ml.Phaseis a string literal with values in["train", "valid", "test"].mlfab.get_data_dir()returns the data directory, which can be set in a configuration file which lives in~/.mlfab.yml. The default configuration file will be written on first run if it doesn't exist yet.defget_dataset(self,phase:mlfab.Phase)->Dataset[tuple[Tensor,Tensor]]:root_dir=mlfab.get_data_dir()/"mnist"returnMNIST(root_dir=root_dir,train=phase=="train")Optimizersdefbuild_optimizer(self)->Optimizer:returnmlfab.Adam.get(self,lr=1e-3)Compute LossEachmlfabmodel should either implement theforwardfunction, which should take a batch from the dataset and return the loss, or, if more control is desired, theget_lossfunction can be overridden.defforward(self,x:Tensor)->Tensor:returnself.model(x)defget_loss(self,batch:tuple[Tensor,Tensor],state:mlfab.State)->Tensor:x,y=batchyhat=self(x)self.log_step(batch,yhat,state)returnF.cross_entropy(yhat,y.squeeze(-1))LoggingWhen we calllog_stepin theget_lossfunction, it delegates to eitherlog_train_step,log_valid_steporlog_test_step, depending on whatstate.phaseis. In this case, on each validation step we log images of the MNIST digits with the labels that our model predicts.deflog_valid_step(self,batch:tuple[Tensor,Tensor],output:Tensor,state:mlfab.State)->None:(x,y),yhat=batch,outputdefget_label_strings()->list[str]:ytrue,ypred=y.squeeze(-1),yhat.argmax(-1)return[f"ytrue={ytrue[i]}, ypred={ypred[i]}"foriinrange(len(ytrue))]self.log_labeled_images("images",lambda:(x,get_label_strings()))RunningWe can launch a training job using thelaunchclass method. The config can be aConfigobject, or it can be the path to aconfig.yamlfile located in the same directory as the task file. You can additionally provide thelauncherargument, which supports training the model across multiple GPUs or nodes.if__name__=="__main__":MnistClassification.launch(Config(batch_size=16))