Search code examples
juliaflux

How to pass sample weights to a mini-batch for training (Julia Flux)

As mentioned in the title, this is part of my code to train the model on a mini-batch for a single step.

    loss, grad = Flux.withgradient(modelDQN.evalParameters) do
        qEval = modelDQN.evalModel(evalInput)
        Flux.mse(qEval, qTarget)
    end
    Flux.update!(modelDQN.optimizer, modelDQN.evalParameters, grad)

Now I wanted to add a sample weights vector ISWeights(equal length to qEval and qTarget) to the training process to achieve the following function: enter image description here

where for each sample I can assign a weight to adjust the total error and so to the gradient. How can I achieve this using Julia Flux? For better understanding, I enclose the fit function of Tensorflow Keras here, I expect ISWeights to work like sample_weight argument here.

Model.fit(
    x=None,
    y=None,
    batch_size=None,
    epochs=1,
    verbose="auto",
    callbacks=None,
    validation_split=0.0,
    validation_data=None,
    shuffle=True,
    class_weight=None,
    sample_weight=None, # equivalent to this argument here
    initial_epoch=0,
    steps_per_epoch=None,
    validation_steps=None,
    validation_batch_size=None,
    validation_freq=1,
    max_queue_size=10,
    workers=1,
    use_multiprocessing=False,
)
Solution

  • I just found that in the official docs there is a standard and concise way to implement the weighted mean for the error using built-in MSE.

    loss(ŷ, y, agg=x->mean(w .* x))

    Reference: Loss Functions Flux