Write Models¶
This section introduces how to implement a new model entirely from scratch and make it compatible with LiBai.
Construct Models in LiBai¶
LiBai uses LazyConfig for a more flexible config system, which means you can simply import your own model in your config and train it under LiBai.
For image classification task, the input data is usually a batch of images and labels. The following code shows how to build a toy model for this task. Import in your code:
# toy_model.py
import oneflow as flow
import oneflow.nn as nn
class ToyModel(nn.Module):
def __init__(self,
num_classes=1000,
):
super().__init__()
self.features = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
self.avgpool = nn.AdaptiveAvgPool2d(1)
self.classifier = nn.Linear(64, num_classes)
self.loss_func = nn.CrossEntropyLoss()
def forward(self, images, labels=None):
x = self.features(images)
x = self.avgpool(x)
x = flow.flatten(x, 1)
x = self.classifier(x)
if labels is not None and self.training:
losses = self.loss_func(x, labels)
return {"losses": losses}
else:
return {"prediction_scores": x}
Note:
For classification models, the
forward
function must haveimages
andlabels
as arguments, which correspond to the output in__getitem__
of LiBai’s built-in datasets. Please refer to imagenet.py for more details about the dataset.This toy model will return
losses
during training andprediction_scores
during inference, and both of them should be the type ofdict
, which means you should implement theloss function
in your model, likeself.loss_func=nn.CrossEntropyLoss()
as the ToyModel shows above.
Import the model in config¶
With LazyConfig System
, you can simply import the model in your config file. The following code shows how to use ToyModel
in your config file:
# config.py
from libai.config import LazyCall
from toy_model import ToyModel
model = LazyCall(ToyModel)(
num_classes=1000
)