I am trying to implement an operator that will change the labels for computing the loss function. The loss function is like this:
class OhemCELoss(nn.Module): def __init__(self, thresh, gnore_lb=255): super(OhemCELoss, self).__init__() self.score_thresh = thresh self.ignore_lb = ignore_lb self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='mean') def forward(self, logits, labels): import ohem_cpp n_min = (labels != self.ignore_lb).numel() // 16 labels = ohem_cpp.score_ohem_label(logits, labels, self.ignore_lb, self.score_thresh, n_min).detach() loss = self.criteria(logits, labels) return loss
ohem_cpp is implemented according to the principle of the pytorch extension, and the
ohem_cpp.score_ohem_label will compute the softmax scores of the input logits, choose the scores lower than the
thresh and set the corresponding element of the
labels to the assigned
ignore_lb. I noticed that in pytorch1.6, users should add decorator of
@custum_fwd if they want to implement their own operators that is derived from
autograd.Function. I have no idea whether I can call my
ohem_cpp module directly as in the code above without potential errors. Did I miss anything here ?
I am asking this because I sometimes met the error of
terminate called after throwing an instance of 'thrust::system::system_error'(not often), I do not know whether is is associated with my wrongly using of my self-implemented operators.