Build recommendation models with EasyRec

Block properties

Each block has a unique name and one or more inputs. Inputs can be a feature group name, another block's name, or a block package name.

• Multiple inputs: If inputs are lists, they merge into a single list. If inputs are tensors, they concatenate along the last dimension.

• Multiple outputs: A block with multiple outputs returns a Python tuple. Use input_slice or input_fn on a downstream block to select specific elements.

• Input layer: Associate an input_layer with a block to apply batch normalization, layer normalization, or feature dropout to the input feature group. When a block has an input_layer, you must configure feature_group_name as the name of a feature group. If a block is not associated with an input_layer, the name of the block cannot be the same as the name of a feature group.

A component built using the backbone outperforms the equivalent built-in model because the MLP layer uses optimal initialization methods.

Wide & Deep model

Configuration file: wide_and_deep_backbone_on_movielens.config

Wide & Deep uses two blocks — one for the wide (linear) path and one for the deep (MLP) path — then combines them with an Add layer.

model_config: {
 model_name: "WideAndDeep"
 model_class: "RankModel"
 feature_groups: {
   group_name: 'wide'
   feature_names: 'user_id'
   feature_names: 'movie_id'
   feature_names: 'job_id'
   feature_names: 'age'
   feature_names: 'gender'
   feature_names: 'year'
   feature_names: 'genres'
   wide_deep: WIDE
 }
 feature_groups: {
   group_name: 'deep'
   feature_names: 'user_id'
   feature_names: 'movie_id'
   feature_names: 'job_id'
   feature_names: 'age'
   feature_names: 'gender'
   feature_names: 'year'
   feature_names: 'genres'
   wide_deep: DEEP
 }
 backbone {
   blocks {
     name: 'wide'
     inputs {
     	feature_group_name: 'wide'
     }
     input_layer {
     	only_output_feature_list: true
     	wide_output_dim: 1
     }
   }
   blocks {
     name: 'deep_logit'
     inputs {
     	feature_group_name: 'deep'
     }
     keras_layer {
     	 class_name: 'MLP'
       mlp {
         hidden_units: [256, 256, 256, 1]
         use_final_bn: false
         final_activation: 'linear'
       }
     }
   }
   blocks {
     name: 'final_logit'
     inputs {
       block_name: 'wide'
       input_fn: 'lambda x: tf.add_n(x)'
     }
     inputs {
       block_name: 'deep_logit'
     }
     merge_inputs_into_list: true
     keras_layer {
       class_name: 'Add'
     }
   }
   concat_blocks: 'final_logit'
 }
 model_params {
 	l2_regularization: 1e-4
 }
 embedding_regularization: 1e-4
}

Performance on the MovieLens-1M dataset:

DeepFM model

Configuration file: deepfm_backbone_on_movielens.config

This example demonstrates two special block types: a block with a custom lambda function and a block that loads the built-in tf.keras.layers.Add layer.

model_config: {
  model_name: 'DeepFM'
  model_class: 'RankModel'
  feature_groups: {
    group_name: 'wide'
    feature_names: 'user_id'
    feature_names: 'movie_id'
    feature_names: 'job_id'
    feature_names: 'age'
    feature_names: 'gender'
    feature_names: 'year'
    feature_names: 'genres'
    wide_deep: WIDE
  }
  feature_groups: {
    group_name: 'features'
    feature_names: 'user_id'
    feature_names: 'movie_id'
    feature_names: 'job_id'
    feature_names: 'age'
    feature_names: 'gender'
    feature_names: 'year'
    feature_names: 'genres'
    feature_names: 'title'
    wide_deep: DEEP
  }
  backbone {
    blocks {
      name: 'wide_logit'
      inputs {
        feature_group_name: 'wide'
      }
      input_layer {
        wide_output_dim: 1
      }
    }
    blocks {
      name: 'features'
      inputs {
        feature_group_name: 'features'
      }
      input_layer {
        output_2d_tensor_and_feature_list: true
      }
    }
    blocks {
      name: 'fm'
      inputs {
        block_name: 'features'
        input_slice: '[1]'
      }
      keras_layer {
        class_name: 'FM'
      }
    }
    blocks {
      name: 'deep'
      inputs {
        block_name: 'features'
        input_slice: '[0]'
      }
      keras_layer {
        class_name: 'MLP'
        mlp {
          hidden_units: [256, 128, 64, 1]
          use_final_bn: false
          final_activation: 'linear'
        }
      }
    }
    blocks {
      name: 'add'
      inputs {
        block_name: 'wide_logit'
        input_fn: 'lambda x: tf.reduce_sum(x, axis=1, keepdims=True)'
      }
      inputs {
        block_name: 'fm'
      }
      inputs {
        block_name: 'deep'
      }
      merge_inputs_into_list: true
      keras_layer {
        class_name: 'Add'
      }
    }
    concat_blocks: 'add'
  }
  model_params {
    l2_regularization: 1e-4
  }
  embedding_regularization: 1e-4
}

Performance on the MovieLens-1M dataset:

DCN model

Configuration file: dcn_backbone_on_movielens.config

This example uses a recurrent block to loop a Cross layer three times, implementing the Deep & Cross Network (DCN) v2 architecture. A top MLP layer is added to the backbone.

model_config: {
  model_name: 'DCN V2'
  model_class: 'RankModel'
  feature_groups: {
    group_name: 'all'
    feature_names: 'user_id'
    feature_names: 'movie_id'
    feature_names: 'job_id'
    feature_names: 'age'
    feature_names: 'gender'
    feature_names: 'year'
    feature_names: 'genres'
    wide_deep: DEEP
  }
  backbone {
    blocks {
      name: "deep"
      inputs {
        feature_group_name: 'all'
      }
      keras_layer {
        class_name: 'MLP'
        mlp {
          hidden_units: [256, 128, 64]
        }
      }
    }
    blocks {
      name: "dcn"
      inputs {
        feature_group_name: 'all'
        input_fn: 'lambda x: [x, x]'
      }
      recurrent {
        num_steps: 3
        fixed_input_index: 0
        keras_layer {
          class_name: 'Cross'
        }
      }
    }
    concat_blocks: ['deep', 'dcn']
    top_mlp {
      hidden_units: [64, 32, 16]
    }
  }
  model_params {
    l2_regularization: 1e-4
  }
  embedding_regularization: 1e-4
}

The cross layer looped 3 times is logically equivalent to:

x1 = Cross()(x0, x0)
x2 = Cross()(x0, x1)
x3 = Cross()(x0, x2)

Performance on the MovieLens-1M dataset:

The backbone version uses the DCN v2 Cross component, which has more parameters than the built-in v1 implementation.

DLRM model

Configuration file: dlrm_backbone_on_criteo.config

This example uses DotInteraction — a component for pairwise feature dot-product interactions. The first input of the dot block is a tensor; the second is a list. The framework inserts the tensor into the list and merges them before passing to the block.

model_config: {
  model_name: 'DLRM'
  model_class: 'RankModel'
  feature_groups: {
    group_name: "dense"
    feature_names: "F1"
    feature_names: "F2"
    ...
    wide_deep:DEEP
  }
  feature_groups: {
    group_name: "sparse"
    feature_names: "C1"
    feature_names: "C2"
    feature_names: "C3"
    ...
    wide_deep:DEEP
  }
  backbone {
    blocks {
      name: 'bottom_mlp'
      inputs {
        feature_group_name: 'dense'
      }
      keras_layer {
        class_name: 'MLP'
        mlp {
          hidden_units: [64, 32, 16]
        }
      }
    }
    blocks {
      name: 'sparse'
      inputs {
        feature_group_name: 'sparse'
      }
      input_layer {
        output_2d_tensor_and_feature_list: true
      }
    }
    blocks {
      name: 'dot'
      inputs {
        block_name: 'bottom_mlp'
      }
      inputs {
        block_name: 'sparse'
        input_slice: '[1]'
      }
      keras_layer {
        class_name: 'DotInteraction'
      }
    }
    blocks {
      name: 'sparse_2d'
      inputs {
        block_name: 'sparse'
        input_slice: '[0]'
      }
    }
    concat_blocks: ['sparse_2d', 'dot']
    top_mlp {
      hidden_units: [256, 128, 64]
    }
  }
  model_params {
    l2_regularization: 1e-5
  }
  embedding_regularization: 1e-5
}

Performance on the Criteo dataset:

DLRM with numerical feature embedding

Configuration file: dlrm_on_criteo_with_periodic.config

Building on the DLRM example, this adds a PeriodicEmbedding layer for numerical features. Parameters are passed using google.protobuf.Struct (the st_params field). The framework also supports custom protobuf messages — both methods share a common parameter API.

model_config: {
  model_class: 'RankModel'
  feature_groups: {
    group_name: "dense"
    feature_names: "F1"
    feature_names: "F2"
    ...
    wide_deep:DEEP
  }
  feature_groups: {
    group_name: "sparse"
    feature_names: "C1"
    feature_names: "C2"
    ...
    wide_deep:DEEP
  }
  backbone {
    blocks {
      name: 'num_emb'
      inputs {
        feature_group_name: 'dense'
      }
      keras_layer {
        class_name: 'PeriodicEmbedding'
        st_params {
          fields {
            key: "output_tensor_list"
            value { bool_value: true }
          }
          fields {
            key: "embedding_dim"
            value { number_value: 16 }
          }
          fields {
            key: "sigma"
            value { number_value: 0.005 }
          }
        }
      }
    }
    blocks {
      name: 'sparse'
      inputs {
        feature_group_name: 'sparse'
      }
      input_layer {
        output_2d_tensor_and_feature_list: true
      }
    }
    blocks {
      name: 'dot'
      inputs {
        block_name: 'num_emb'
        input_slice: '[1]'
      }
      inputs {
        block_name: 'sparse'
        input_slice: '[1]'
      }
      keras_layer {
        class_name: 'DotInteraction'
      }
    }
    blocks {
      name: 'sparse_2d'
      inputs {
        block_name: 'sparse'
        input_slice: '[0]'
      }
    }
    blocks {
      name: 'num_emb_2d'
      inputs {
        block_name: 'num_emb'
        input_slice: '[0]'
      }
    }
    concat_blocks: ['num_emb_2d', 'dot', 'sparse_2d']
    top_mlp {
      hidden_units: [256, 128, 64]
    }
  }
  model_params {
    l2_regularization: 1e-5
  }
  embedding_regularization: 1e-5
}

Performance on the Criteo dataset:

DNN model using built-in Keras layers

Configuration file: mlp_on_movielens.config

This example shows how to use a sequential block — multiple Keras layers connected in series where each layer's output becomes the next layer's input. Built-in Keras layers require parameters in google.protobuf.Struct format.

model_config: {
  model_class: "RankModel"
  feature_groups: {
    group_name: 'features'
    feature_names: 'user_id'
    feature_names: 'movie_id'
    feature_names: 'job_id'
    feature_names: 'age'
    feature_names: 'gender'
    feature_names: 'year'
    feature_names: 'genres'
    wide_deep: DEEP
  }
  backbone {
    blocks {
      name: 'mlp'
      inputs {
        feature_group_name: 'features'
      }
      layers {
        keras_layer {
          class_name: 'Dense'
          st_params {
            fields {
              key: 'units'
              value: { number_value: 256 }
            }
            fields {
              key: 'activation'
              value: { string_value: 'relu' }
            }
          }
        }
      }
      layers {
        keras_layer {
          class_name: 'Dropout'
          st_params {
            fields {
              key: 'rate'
              value: { number_value: 0.5 }
            }
          }
        }
      }
      layers {
        keras_layer {
          class_name: 'Dense'
          st_params {
            fields {
              key: 'units'
              value: { number_value: 256 }
            }
            fields {
              key: 'activation'
              value: { string_value: 'relu' }
            }
          }
        }
      }
      layers {
        keras_layer {
          class_name: 'Dropout'
          st_params {
            fields {
              key: 'rate'
              value: { number_value: 0.5 }
            }
          }
        }
      }
      layers {
        keras_layer {
          class_name: 'Dense'
          st_params {
            fields {
              key: 'units'
              value: { number_value: 1 }
            }
          }
        }
      }
    }
    concat_blocks: 'mlp'
  }
  model_params {
    l2_regularization: 1e-4
  }
  embedding_regularization: 1e-4
}

Performance on the MovieLens-1M dataset:

Contrastive learning with block packages

Configuration file: contrastive_learning_on_movielens.config

A block package encapsulates a set of blocks into a reusable subnetwork that can be called multiple times with shared parameters. This differs from a plain block, which cannot be called multiple times (though its results can be reused). Block packages are designed for self-supervised learning and contrastive learning scenarios.

model_config: {
  model_name: "ContrastiveLearning"
  model_class: "RankModel"
  feature_groups: {
    group_name: 'user'
    feature_names: 'user_id'
    feature_names: 'job_id'
    feature_names: 'age'
    feature_names: 'gender'
    wide_deep: DEEP
  }
  feature_groups: {
    group_name: 'item'
    feature_names: 'movie_id'
    feature_names: 'year'
    feature_names: 'genres'
    wide_deep: DEEP
  }
  backbone {
    blocks {
      name: 'user_tower'
      inputs {
        feature_group_name: 'user'
      }
      keras_layer {
        class_name: 'MLP'
        mlp {
          hidden_units: [256, 128]
        }
      }
    }
    packages {
      name: 'item_tower'
      blocks {
        name: 'item'
        inputs {
          feature_group_name: 'item'
        }
        input_layer {
          dropout_rate: 0.2
        }
      }
      blocks {
        name: 'item_encoder'
        inputs {
          block_name: 'item'
        }
        keras_layer {
          class_name: 'MLP'
          mlp {
            hidden_units: [256, 128]
          }
        }
      }
    }
    blocks {
      name: 'contrastive_learning'
      inputs {
        package_name: 'item_tower'
      }
      inputs {
        package_name: 'item_tower'
      }
      merge_inputs_into_list: true
      keras_layer {
        class_name: 'AuxiliaryLoss'
        st_params {
          fields {
            key: 'loss_type'
            value: { string_value: 'info_nce' }
          }
          fields {
            key: 'loss_weight'
            value: { number_value: 0.1 }
          }
          fields {
            key: 'temperature'
            value: { number_value: 0.2 }
          }
        }
      }
    }
    blocks {
      name: 'top_mlp'
      inputs {
        block_name: 'contrastive_learning'
        ignore_input: true
      }
      inputs {
        block_name: 'user_tower'
      }
      inputs {
        package_name: 'item_tower'
        reset_input {}
      }
      keras_layer {
        class_name: 'MLP'
        mlp {
          hidden_units: [128, 64]
        }
      }
    }
    concat_blocks: 'top_mlp'
  }
  model_params {
    l2_regularization: 1e-4
  }
  embedding_regularization: 1e-4
}

In this example, item_tower is called three times:

• First two calls: The dropout_rate: 0.2 input layer configuration takes effect, generating augmented samples for the contrastive learning loss.

• Third call: reset_input {} overrides the input layer configuration, disabling dropout for the main model prediction path.

The item_tower package shares parameters across all three calls. The contrastive auxiliary task and the main model use the same encoder weights.

Two special input configurations used in this example:

• ignore_input: true — the input is not used for computation but controls DAG execution order (the contrastive_learning block runs before top_mlp).

• reset_input — overrides the input layer parameters defined in the package for that specific call.

AuxiliaryLoss calculates the contrastive learning loss. For parameter details, see Component parameters.

No concat_blocks is set for the item_tower package, so the framework automatically treats it as a DAG leaf node.

Performance on the MovieLens-1M dataset:

For a more complex contrastive learning example, see CL4SRec.

Multi-objective model: MMoE

For multi-objective models, set model_class to MultiTaskModel and configure multiple task towers in model_params. The model_name field is a free-form string used only for comments.

model_config {
  model_name: "MMoE"
  model_class: "MultiTaskModel"
  feature_groups {
    group_name: "all"
    feature_names: "user_id"
    feature_names: "cms_segid"
    ...
    feature_names: "tag_brand_list"
    wide_deep: DEEP
  }
  backbone {
    blocks {
      name: 'all'
      inputs {
        feature_group_name: 'all'
      }
      input_layer {
        only_output_feature_list: true
      }
    }
    blocks {
      name: "senet"
      inputs {
        block_name: "all"
      }
      keras_layer {
        class_name: 'SENet'
        senet {
          reduction_ratio: 4
        }
      }
    }
    blocks {
      name: "mmoe"
      inputs {
        block_name: "senet"
      }
      keras_layer {
        class_name: 'MMoE'
        mmoe {
          num_task: 2
          num_expert: 3
          expert_mlp {
            hidden_units: [256, 128]
          }
        }
      }
    }
  }
  model_params {
    task_towers {
      tower_name: "ctr"
      label_name: "clk"
      dnn {
        hidden_units: [128, 64]
      }
      num_class: 1
      weight: 1.0
      loss_type: CLASSIFICATION
      metrics_set: {
       auc {}
      }
    }
    task_towers {
      tower_name: "cvr"
      label_name: "buy"
      dnn {
        hidden_units: [128, 64]
      }
      num_class: 1
      weight: 1.0
      loss_type: CLASSIFICATION
      metrics_set: {
       auc {}
      }
    }
    l2_regularization: 1e-06
  }
  embedding_regularization: 5e-05
}

No concat_blocks is set for the backbone, so the framework automatically treats it as a DAG leaf node.

Multi-objective model: DBMTL

Deep Bayesian Multi-Task Learning (DBMTL) requires relation_dnn for each task tower and uses relation_tower_names to configure inter-task dependencies.

model_config {
 model_name: "DBMTL"
 model_class: "MultiTaskModel"
 feature_groups {
 group_name: "all"
 feature_names: "user_id"
 feature_names: "cms_segid"
 ...
 feature_names: "tag_brand_list"
 wide_deep: DEEP
 }
 backbone {
 blocks {
 name: "mask_net"
 inputs {
 feature_group_name: "all"
 }
 keras_layer {
 class_name: 'MaskNet'
 masknet {
 mask_blocks {
 aggregation_size: 512
 output_size: 256
 }
 mask_blocks {
 aggregation_size: 512
 output_size: 256
 }
 mask_blocks {
 aggregation_size: 512
 output_size: 256
 }
 mlp {
 hidden_units: [512, 256]
 }
 }
 }
 }
 }
 model_params {
 task_towers {
 tower_name: "ctr"
 label_name: "clk"
 loss_type: CLASSIFICATION
 metrics_set: {
 auc {}
 }
 dnn {
 hidden_units: [256, 128, 64]
 }
 relation_dnn {
 hidden_units: [32]
 }
 weight: 1.0
 }
 task_towers {
 tower_name: "cvr"
 label_name: "buy"
 loss_type: CLASSIFICATION
 metrics_set: {
 auc {}
 }
 dnn {
 hidden_units: [256, 128, 64]
 }
 relation_tower_names: ["ctr"]
 relation_dnn {
 hidden_units: [32]
 }
 weight: 1.0
 }
 l2_regularization: 1e-6
 }
 embedding_regularization: 5e-6
}

No concat_blocks is set for the backbone, so the framework automatically treats it as a DAG leaf node.

MaskNet + PPNet + MMoE

This example demonstrates a reusable block combining three components in sequence.

model_config: {
  model_name: 'MaskNet + PPNet + MMoE'
  model_class: 'RankModel'
  feature_groups: {
    group_name: 'memorize'
    feature_names: 'user_id'
    feature_names: 'adgroup_id'
    feature_names: 'pid'
    wide_deep: DEEP
  }
  feature_groups: {
    group_name: 'general'
    feature_names: 'age_level'
    feature_names: 'shopping_level'
    ...
    wide_deep: DEEP
  }
  backbone {
    blocks {
      name: "mask_net"
      inputs {
        feature_group_name: "general"
      }
      repeat {
        num_repeat: 3
        keras_layer {
          class_name: "MaskBlock"
          mask_block {
            output_size: 512
            aggregation_size: 1024
          }
        }
      }
    }
    blocks {
      name: "ppnet"
      inputs {
        block_name: "mask_net"
      }
      inputs {
        feature_group_name: "memorize"
      }
      merge_inputs_into_list: true
      repeat {
        num_repeat: 3
        input_fn: "lambda x, i: [x[0][i], x[1]]"
        keras_layer {
          class_name: "PPNet"
          ppnet {
            mlp {
              hidden_units: [256, 128, 64]
            }
            gate_params {
              output_dim: 512
            }
            mode: "eager"
            full_gate_input: false
          }
        }
      }
    }
    blocks {
      name: "mmoe"
      inputs {
        block_name: "ppnet"
      }
      inputs {
        feature_group_name: "general"
      }
      keras_layer {
        class_name: "MMoE"
        mmoe {
          num_task: 2
          num_expert: 3
        }
      }
    }
  }
  model_params {
    l2_regularization: 0.0
    task_towers {
      tower_name: "ctr"
      label_name: "is_click"
      metrics_set {
        auc {
          num_thresholds: 20000
        }
      }
      loss_type: CLASSIFICATION
      num_class: 1
      dnn {
        hidden_units: 64
        hidden_units: 32
      }
      weight: 1.0
    }
    task_towers {
      tower_name: "cvr"
      label_name: "is_train"
      metrics_set {
        auc {
          num_thresholds: 20000
        }
      }
      loss_type: CLASSIFICATION
      num_class: 1
      dnn {
        hidden_units: 64
        hidden_units: 32
      }
      weight: 1.0
    }
  }
}

More examples

New models:

• FiBiNet: fibinet_on_movielens.config

• MaskNet: masknet_on_movielens.config

Performance on the MovieLens-1M dataset:

Sequential models:

• Deep Interest Network (DIN): DIN_backbone.config

• Behavior Sequence Transformer (BST): BST_backbone.config

• CL4SRec: CL4SRec

Other models:

• Highway Network: Highway Network

• Cross Decoupling Network (CDN): CDN

• DLRM+SENet: dlrm_senet_on_criteo.config

Basic components

Feature crossing components

Feature importance learning components

Sequential feature encoding components

Multi-objective learning components

Auxiliary loss function components

For component parameter details, see Component parameters.

The preceding reference is in Chinese.

Empty block

An empty block has no layer configured and is used only to merge multiple inputs.

Input block

An input block is associated with an input_layer to retrieve, process, and return raw feature inputs. Input blocks accept only one input, specified by feature_group_name. The block name must match the feature group name.

blocks {
 name: 'all'
 inputs {
 feature_group_name: 'all'
 }
 input_layer {
 only_output_feature_list: true
 }
}

The InputLayer protobuf parameters:

message InputLayer {
 optional bool do_batch_norm = 1;
 optional bool do_layer_norm = 2;
 optional float dropout_rate = 3;
 optional float feature_dropout_rate = 4;
 optional bool only_output_feature_list = 5;
 optional bool only_output_3d_tensor = 6;
 optional bool output_2d_tensor_and_feature_list = 7;
 optional bool output_seq_and_normal_feature = 8;
}

Lambda block

A Lambda block applies a lambda expression for simple transformations.

blocks {
 name: 'wide_logit'
 inputs {
 feature_group_name: 'wide'
 }
 lambda {
 expression: 'lambda x: tf.reduce_sum(x, axis=1, keepdims=True)'
 }
}

KerasLayer block

A KerasLayer block loads and executes a Keras layer — either a custom layer or a built-in TensorFlow Keras layer.

• class_name: The Keras layer class name to load.

• st_params: Parameters in google.protobuf.Struct format (required for built-in Keras layers).

• Custom protobuf message parameters are also supported.

keras_layer {
 class_name: 'MLP'
 mlp {
 hidden_units: [64, 32, 16]
 }
}

keras_layer {
 class_name: 'Dropout'
 st_params {
 fields {
 key: 'rate'
 value: { number_value: 0.5 }
 }
 }
}

Recurrent block

A recurrent block implements an RNN-like loop, executing a layer multiple times. Each execution receives the output of the previous one.

recurrent {
 num_steps: 3
 fixed_input_index: 0
 keras_layer {
 class_name: 'Cross'
 }
}

For a usage example, see DCN model.

Repeated block

A repeated block executes a component multiple times with the same inputs, implementing multi-head logic.

repeat {
 num_repeat: 2
 keras_layer {
 class_name: "MaskBlock"
 mask_block {
 output_size: 512
 aggregation_size: 2048
 input_layer_norm: false
 }
 }
}

For a usage example, see MaskNet + PPNet + MMoE.

Sequential block

A sequential block executes multiple layers in series, where each layer's output becomes the next layer's input. This is more concise than chaining multiple common blocks end-to-end.

blocks {
 name: 'mlp'
 inputs {
 feature_group_name: 'features'
 }
 layers {
 keras_layer {
 class_name: 'Dense'
 st_params {
 fields {
 key: 'units'
 value: { number_value: 256 }
 }

 fields {
 key: 'activation'
 value: { string_value: 'relu' }
 }
 }
 }
 }
 layers {
 keras_layer {
 class_name: 'Dropout'
 st_params {
 fields {
 key: 'rate'
 value: { number_value: 0.5 }
 }
 }
 }
 }
 layers {
 keras_layer {
 class_name: 'Dense'
 st_params {
 fields {
 key: 'units'
 value: { number_value: 1 }
 }
 }
 }
 }
}

Block packages

A block package encapsulates a DAG of multiple blocks and can be called multiple times with shared parameters. Block packages are used in self-supervised learning and contrastive learning scenarios.

The protobuf definition:

message BlockPackage {
 // package name
 required string name = 1;
 // a few blocks generating a DAG
 repeated Block blocks = 2;
 // the names of output blocks
 repeated string concat_blocks = 3;
}

Reference a block package in a block using package_name. For a complete usage example, see Contrastive learning with block packages.

The following example shows how to use a block package to implement contrastive learning with a BST-based encoder:

model_config {
 model_class: "RankModel"
 feature_groups {
 group_name: "all"
 feature_names: "adgroup_id"
 feature_names: "user"
 ...
 feature_names: "pid"
 wide_deep: DEEP
 }

 backbone {
 packages {
 name: 'feature_encoder'
 blocks {
 name: "fea_dropout"
 inputs {
 feature_group_name: "all"
 }
 input_layer {
 dropout_rate: 0.5
 only_output_3d_tensor: true
 }
 }
 blocks {
 name: "encode"
 inputs {
 block_name: "fea_dropout"
 }
 layers {
 keras_layer {
 class_name: 'BSTCTR'
 bst {
 hidden_size: 128
 num_attention_heads: 4
 num_hidden_layers: 3
 intermediate_size: 128
 hidden_act: 'gelu'
 max_position_embeddings: 50
 hidden_dropout_prob: 0.1
 attention_probs_dropout_prob: 0
 }
 }
 }
 layers {
 keras_layer {
 class_name: 'Dense'
 st_params {
 fields {
 key: 'units'
 value: { number_value: 128 }
 }
 fields {
 key: 'kernel_initializer'
 value: { string_value: 'zeros' }
 }
 }
 }
 }
 }
 }
 blocks {
 name: "all"
 inputs {
 name: "all"
 }
 input_layer {
 only_output_3d_tensor: true
 }
 }
 blocks {
 name: "loss_ctr"
 merge_inputs_into_list: true
 inputs {
 package_name: 'feature_encoder'
 }
 inputs {
 package_name: 'feature_encoder'
 }
 inputs {
 package_name: 'all'
 }
 keras_layer {
 class_name: 'LOSSCTR'
 st_params{
 fields {
 key: 'cl_weight'
 value: { number_value: 1 }
 }
 fields {
 key: 'au_weight'
 value: { number_value: 0.01 }
 }
 }
 }
 }
 }
 model_params {
 l2_regularization: 1e-5
 }
 embedding_regularization: 1e-5
}

Define the component class

Define a class that inherits tf.keras.layers.Layer. Implement at least two methods: __init__ and call.

def __init__(self, params, name='xxx', reuse=None, **kwargs):
    pass

def call(self, inputs, training=None, **kwargs):
    pass

Implement __init__

The params parameter receives all configuration passed by the framework. It supports both google.protobuf.Struct and custom protobuf message formats through a unified API:

The reuse parameter controls weight reuse. Declare all dependent Keras layers in __init__ using the native tf.keras.layers.* API. You can also use the tf.layers.* function and pass the reuse parameter based on your business requirements.

Implement call

The call method contains the main component logic. The inputs parameter is a tensor or a list of tensors. The optional training parameter indicates whether the model is in training mode.

Register the component

Export the new layer from easy_rec.python.layers.keras.__init__.py so the framework can recognize it as a component library member. For example, to export the MLP class from blocks.py:

from .blocks import MLP

(Optional) Add a custom protobuf message

To pass parameters via a custom protobuf message:

1. Add the message definition to easy_rec/python/protos/layer.proto.

2. Register the parameter in the KerasLayer.params message body in easy_rec/python/protos/keras_layer.proto.

Example: FM layer

class FM(tf.keras.layers.Layer):
    """Factorization Machine models pairwise (order-2) feature interactions without linear term and bias.

    References
    - [Factorization Machines](https://www.csie.ntu.edu.tw/~b97053/paper/Rendle2010FM.pdf)
    Input shape.
    - List of 2D tensor with shape: ``(batch_size,embedding_size)``.
    - Or a 3D tensor with shape: ``(batch_size,field_size,embedding_size)``
    Output shape
    - 2D tensor with shape: ``(batch_size, 1)``.
    """

    def __init__(self, params, name='fm', reuse=None, **kwargs):
        super(FM, self).__init__(name, **kwargs)
        self.reuse = reuse
        self.use_variant = params.get_or_default('use_variant', False)

    def call(self, inputs, **kwargs):
        if type(inputs) == list:
            emb_dims = set(map(lambda x: int(x.shape[-1]), inputs))
            if len(emb_dims) != 1:
                dims = ','.join([str(d) for d in emb_dims])
                raise ValueError('all embedding dim must be equal in FM layer:' + dims)
            with tf.name_scope(self.name):
                fea = tf.stack(inputs, axis=1)
        else:
            assert inputs.shape.ndims == 3, 'input of FM layer must be a 3D tensor or a list of 2D tensors'
            fea = inputs

        with tf.name_scope(self.name):
            square_of_sum = tf.square(tf.reduce_sum(fea, axis=1))
            sum_of_square = tf.reduce_sum(tf.square(fea), axis=1)
            cross_term = tf.subtract(square_of_sum, sum_of_square)
            if self.use_variant:
                cross_term = 0.5 * cross_term
            else:
                cross_term = 0.5 * tf.reduce_sum(cross_term, axis=-1, keepdims=True)
            return cross_term