We have a model with an output target that is a 2D tensor. That is, the output represents a set of n classes, evaluated at m bins within the data. That is the output shape is: [none, n, m]. Furthermore, the classes are highly imbalanced, so we need to use weights to balance the losses across the classes.

For example, the target shape with three classes (A, B, and C), four bins and a batch size of 1 would be:

Target = 
 [[[A1, A2, A3, A4],
   [B1, B2, B3, B4],
   [C1, C2, C3, C4]]]

I've researched the use of Class-weights with Keras and it is extensively discussed, but 99% of the time the output is a class vector with a single target result for each class. In our case, we have a target array for each class.

I would hypothesize that we could use a possible weight tensor such as this:

Weights = 
  [[Wa, Wa, Wa, Wa],
   [Wb, Wb, Wb, Wb],
   [Wc, Wc, Wc, Wc]]

Where a constant weight value is applied to each class for each of its bins. But I'm unsure if the shape should be [n, m] or [none, n, m] or [b, n, m] (where b is the batch size). I'm also unsure if this is considered a loss_weight or a class_weight.

I've looked at using loss_weights, class_weights and weight_metrics but the documentation is thin for non-vector outputs.

My question: how does one apply weights to a output tensor?

EDIT 1: I'm starting to think that the samples_weights option may be the best approach here.
Here is a discussion by Francois Chollett offering his thoughts.

  • $\begingroup$ Not sure I fully understand your setting, but the clearest way would be to just write your own loss function $\endgroup$
    – Adam
    Apr 24, 2023 at 22:07

1 Answer 1


Following the recommendation from @Adam we went ahead and built a custom loss function to accept sample-weights. Because we are using a dataset (tf.data) pipeline, we append the sample-weights tensor to the training dataset only, resulting in a three-tuple of: (InputTensor, TargetTensor, WeightTensor). For the test/val dataset, we do not append the weights, resulting in a two-tuple of: (InputTensor, TargetTensor).
The sample-weights tensor is constructed as listed above in the question:

Weights = 
  [[Wa, Wa, Wa, Wa],
   [Wb, Wb, Wb, Wb],
   [Wc, Wc, Wc, Wc]]

The loss function is:

class WeightedCategoricalCrossentropySimple(tf.keras.losses.Loss):
    def __init__(self, y_true, y_pred, **kwargs):   # Note: these variables are NOT used, but there is a bug in TF that calls the constructor incorrectly with these values; we need to consume them but dont need to use them

    def call(self, y_true, y_pred, sample_weight = None):
        # Compute cross-entropy loss for each sample in batch
        loss_per_sample = tf.keras.losses.categorical_crossentropy(y_true, y_pred)

        # Apply weighting to each sample in batch, but only if present
        if sample_weight:
            loss_per_sample = tf.multiply(loss_per_sample, sample_weight)

        loss = tf.reduce_mean(loss_per_sample)

        return loss

The dataset pipeline creation code adds one line to apply the weights to the pipelines (if present) just before batching the tensors:

def load_tfrecord_dataset(...
                          weights = None             # type: np.array
    dataset = tf.data.Dataset.list_files(file_pattern)
    dataset = dataset.interleave(
        lambda tfr: tf.data.TFRecordDataset(tfr, compression_type="GZIP"),

    dataset = dataset.shuffle(buffer_size=buffer_size).repeat(1)
    dataset = dataset.map(_parse_tfrecord_single_class, num_parallel_calls=5)

    # Appends the weight matrix as a third parameter, but only if present
    dataset = dataset.map(lambda x, y: append_weights(x, y, weights))

    dataset = dataset.batch(batch_size=batch_size, num_parallel_calls=tf.data.AUTOTUNE, drop_remainder=True)
    dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    return dataset

Where append_weights() is defined as:

def append_weights(x, y, weights):
    if weights is not None:
        return x, y, weights
        return x, y

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.