If you want to design a new model, you need to write a class inherited from BaseModel
, e.g, NewModel in newmodel.py:
from ..base_model import BaseModel
class NewModel(BaseModel):
def __init__(self, cfgs, is_training):
super().__init__(cfgs, is_training)
def build_network(self, model_cfg):
self.encoder = ...
def forward(self, inputs):
ipts, labs, typs, viws, seqL = inputs
pose = ipts[0]
if len(pose.size()) == 4:
pose = pose.unsqueeze(2)
del ipts
n, s, c, h, w = pose.size()
embed_1, logits, embed = self.encoder(pose)
return {
'training_feat': {
'triplet': {'embeddings': embed_1, 'labels': labs},
'softmax': {'logits': logits, 'labels': labs}
},
'visual_summary': {
'image/pose': pose.view(n*s, 1, h, w)
},
'inference_feat': {
'embeddings': embed
}
}
In your model class, at least you need to implement build_network()
and forward()
functions. The first is used to build the networks, and it does not need return value
. Another is used to calculate the features, the return value
is fixed in dictionary format
training_feat
is for the loss computing, and it must be adict
object.
visual_summary
is for visualization, and it must be adict
object.
inference_feat
is for the inference, and it must be adict
object.
triplet
andsoftmax
are the prefixes (or names) of the loss function.
embeddings
,logits
andlabels
are the input arguments of the loss function.
More information can be seen in base_model.py and loss_aggregator.py.
After finishing the model file, you have two steps left to do:
Step 1: Put your newmodel.py under fastposegait/modeling/models
.
Step 2: Specify the model name in a yaml file:
model_cfg:
model: NewModel
param1: ...
param2: ...
param3: ...
If you want to write a new loss, you need to write a class inherited from fastposegait/modeling/losses
, like this
from .base import BaseLoss
class NewLoss(BaseLoss):
def __init__(self, *args, **kwargs):
super(NewLoss, self).__init__(*args, **kargs)
@gather_and_scale_wrapper
def forward(self, embeddings, labels):
pass
Remember to use gather_and_scale_wrapper
to wrap your forward function if your loss is computed by pairs like triplet
. By this, we gather all features to one GPU and scale the loss by the number of GPUs.
Then, put your loss in fastposegait/modeling/losses
so that you can use it in config file.
Moreover, refer to loss_aggregator.py to explore how your defined loss works in the model.