class_docstr: Defines commonalities between all semantic segmentation models.
  
  A model is class with three members: get_metrics_fn, loss_fn, and a
  flax_model.
  
  get_metrics_fn returns a callable function, metric_fn, that calculates the
  metrics and returns a dictionary. The metric function computes f(x_i, y_i) on
  a minibatch, it has API:
    ```metric_fn(logits, label, weights).```
  
  The trainer will then aggregate and compute the mean across all samples
  evaluated.
  
  loss_fn is a function of API
    loss = loss_fn(logits, batch, model_params=None).
  
  This model class defines a softmax_cross_entropy_loss with weight decay,
  where the weight decay factor is determined by config.l2_decay_factor.
  
  flax_model is returned from the build_flax_model function. A typical
  usage pattern will be:
    ```
    model_cls = model_lib.models.get_model_cls('simple_cnn_segmentation')
    model = model_cls(config, dataset.meta_data)
    flax_model = model.build_flax_model
    dummy_input = jnp.zeros(input_shape, model_input_dtype)
    model_state, params = flax_model.init(
        rng, dummy_input, train=False).pop('params')
    ```
  And this is how to call the model:
    variables = {'params': params, **model_state}
    logits, new_model_state = flax_model.apply(variables, inputs, ...)
    ```
class_docstr_tok: ["Defines","commonalities","between","all","semantic","segmentation","models",".","A","model","is","class","with","three","members",":","get_metrics_fn",",","loss_fn",",","and","a","flax_model",".","get_metrics_fn","returns","a","callable","function",",","metric_fn",",","that","calculates","the","metrics","and","returns","a","dictionary",".","The","metric","function","computes","f","(","x_i",",","y_i",")","on","a","minibatch",",","it","has","API",":","``","`","metric_fn","(","logits",",","label",",","weights",")",".","``","`","The","trainer","will","then","aggregate","and","compute","the","mean","across","all","samples","evaluated",".","loss_fn","is","a","function","of","API","loss","=","loss_fn","(","logits",",","batch",",","model_params=None",")",".","This","model","class","defines","a","softmax_cross_entropy_loss","with","weight","decay",",","where","the","weight","decay","factor","is","determined","by","config.l2_decay_factor",".","flax_model","is","returned","from","the","build_flax_model","function",".","A","typical","usage","pattern","will","be",":","``","`","model_cls","=","model_lib.models.get_model_cls","(","'simple_cnn_segmentation","'",")","model","=","model_cls","(","config",",","dataset.meta_data",")","flax_model","=","model.build_flax_model","dummy_input","=","jnp.zeros","(","input_shape",",","model_input_dtype",")","model_state",",","params","=","flax_model.init","(","rng",",","dummy_input",",","train=False",")",".pop","(","'params","'",")","``","`","And","this","is","how","to","call","the","model",":","variables","=","{","'params","'",":","params",",","*","*","model_state","}","logits",",","new_model_state","=","flax_model.apply","(","variables",",","inputs",",","...",")","``","`"]
docstr: Returns labels' weights to be used for computing weighted loss.
  
  This can used for weighting the loss terms based on the amount of available
  data for each class, when we have un-balances data for different classes.
docstr_tok: ["Returns","labels","'","weights","to","be","used","for","computing","weighted","loss",".","This","can","used","for","weighting","the","loss","terms","based","on","the","amount","of","available","data","for","each","class",",","when","we","have","un-balances","data","for","different","classes","."]
code: def get_label_weights(self) -> jnp.ndarray:
      """Returns labels' weights to be used for computing weighted loss.
  
      This can used for weighting the loss terms based on the amount of available
      data for each class, when we have un-balances data for different classes.
      """
      if not self.config.get("class_rebalancing_factor"):
          return None  # pytype: disable=bad-return-type  # jax-ndarray
      if "class_proportions" not in self.dataset_meta_data:
          raise ValueError(
              "When `class_rebalancing_factor` is nonzero, `class_proportions` must"
              " be provided in `dataset_meta_data`."
          )
      w = self.config.get("class_rebalancing_factor")
      assert (
          0.0 <= w <= 1.0
      ), "`class_rebalancing_factor` must be in [0.0, 1.0]"
      proportions = self.dataset_meta_data["class_proportions"]
      proportions = np.maximum(proportions / np.sum(proportions), 1e-8)
      # Interpolate between no rebalancing (w==0.0) and full reweighting (w==1.0):
      proportions = w * proportions + (1.0 - w)
      weights = 1.0 / proportions
      weights /= np.sum(weights)  # Normalize so weights sum to 1.
      weights *= len(weights)  # Scale so weights sum to num_classes.
      return weights
  
file_path: scenic/scenic/model_lib/base_models/segmentation_model.py
imports: import functools
  from typing import Any, Callable, List, Dict, Optional, Tuple, Union
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 3 for calling the SegmentationModel obj's underlying member method code with example usage: obj.get_label_weights() and returns: weights, None