To use sample_weight, pass it as part of a dictionary for X. This ensures weights are correctly split into train/valid sets and batched. You must also:
- Ensure your
forward method accepts the weight argument. - Set
criterion__reduce=False so the loss is not reduced before weighting. - Override
get_loss to apply the weights to the unreduced loss.
X, y = get_data()
# put your X into a dict if not already a dict
X = {'data': X}
# add sample_weight to the X dict
X['sample_weight'] = sample_weight
class MyModule(nn.Module):
...
def forward(self, data, sample_weight):
# when X is a dict, its keys are passed as kwargs to forward;
# usually, sample_weight can be ignored here
...
class MyNet(NeuralNet):
def __init__(self, *args, criterion__reduce=False, **kwargs):
# make sure to set reduce=False in your criterion, since we need the loss
# for each sample so that it can be weighted
super().__init__(*args, criterion__reduce=criterion__reduce, **kwargs)
def get_loss(self, y_pred, y_true, X, *args, **kwargs):
# override get_loss to use the sample_weight from X
loss_unreduced = super().get_loss(y_pred, y_true, X, *args, **kwargs)
sample_weight = skorch.utils.to_tensor(X['sample_weight'], device=self.device)
loss_reduced = (sample_weight * loss_unreduced).mean()
return loss_reduced
net = MyNet(MyModule, ...)
net.fit(X, y)
