Optional shared encoder for the two tower model. (#1705)

Summary:
Pull Request resolved: #1705

Provide shared encoder option for the two tower model.

Reviewed By: sinonwang

Differential Revision: D29557881

fbshipit-source-id: cd2702c157248fbbfa365358fd23aba075509336

pytext/config/config_adapter.py

@@ -1019,6 +1019,20 @@ def v34_to_v35(json_config):
     return json_config
 
 
+@register_down_grade_adapter(from_version=36)
+def v36_to_v35(json_config):
+    for v in get_json_config_iterator(json_config, "TwoTowerClassificationModel"):
+        if "use_shared_encoder" in v:
+            del v["use_shared_encoder"]
+    return json_config
+
+
+@register_adapter(from_version=35)
+def v35_to_v36(json_config):
+    # New config field was added with backwards-compatible default value
+    return json_config
+
+
 def get_name_from_options(export_config):
     """
     Reverse engineer which model is which based on recognized

pytext/config/pytext_config.py

@@ -409,4 +409,4 @@ class LogitsConfig(TestConfig):
     fp16: bool = False
 
 
-LATEST_VERSION = 35
+LATEST_VERSION = 36

pytext/config/test/json_config/v36_test_downgrade.json

@@ -0,0 +1 @@
+[]

pytext/config/test/json_config/v36_test_upgrade.json

@@ -0,0 +1 @@
+[]

pytext/models/two_tower_classification_model.py

@@ -46,6 +46,7 @@ class InputConfig(ConfigBase):
         output_layer: ClassificationOutputLayer.Config = (
             ClassificationOutputLayer.Config()
         )
+        use_shared_encoder: Optional[bool] = False
 
     def trace(self, inputs):
         return torch.jit.trace(self, inputs)
@@ -113,16 +114,25 @@ def forward(
         left_encoder_inputs: Tuple[torch.Tensor, ...],
         *args
     ) -> List[torch.Tensor]:
-        if self.right_encoder.output_encoded_layers:
-            # if encoded layers are returned, discard them
-            right_representation = self.right_encoder(right_encoder_inputs)[1]
+        if self.use_shared_encoder:
+            if self.right_encoder.output_encoded_layers:
+                # if encoded layers are returned, discard them
+                right_representation = self.right_encoder(right_encoder_inputs)[1]
+                left_representation = self.right_encoder(left_encoder_inputs)[1]
+            else:
+                right_representation = self.right_encoder(right_encoder_inputs)[0]
+                left_representation = self.right_encoder(left_encoder_inputs)[0]
         else:
-            right_representation = self.right_encoder(right_encoder_inputs)[0]
-        if self.left_encoder.output_encoded_layers:
-            # if encoded layers are returned, discard them
-            left_representation = self.left_encoder(left_encoder_inputs)[1]
-        else:
-            left_representation = self.left_encoder(left_encoder_inputs)[0]
+            if self.right_encoder.output_encoded_layers:
+                # if encoded layers are returned, discard them
+                right_representation = self.right_encoder(right_encoder_inputs)[1]
+            else:
+                right_representation = self.right_encoder(right_encoder_inputs)[0]
+            if self.left_encoder.output_encoded_layers:
+                # if encoded layers are returned, discard them
+                left_representation = self.left_encoder(left_encoder_inputs)[1]
+            else:
+                left_representation = self.left_encoder(left_encoder_inputs)[0]
         return self.decoder(right_representation, left_representation, *args)
 
     def caffe2_export(self, tensorizers, tensor_dict, path, export_onnx_path=None):
@@ -178,16 +188,32 @@ def from_config(cls, config: Config, tensorizers: Dict[str, Tensorizer]):
             output_layer_cls = MulticlassOutputLayer
 
         output_layer = output_layer_cls(list(labels), loss)
-        return cls(right_encoder, left_encoder, decoder, output_layer)
+        return cls(
+            right_encoder,
+            left_encoder,
+            decoder,
+            output_layer,
+            config.use_shared_encoder,
+        )
 
     def __init__(
-        self, right_encoder, left_encoder, decoder, output_layer, stage=Stage.TRAIN
+        self,
+        right_encoder,
+        left_encoder,
+        decoder,
+        output_layer,
+        use_shared_encoder,
+        stage=Stage.TRAIN,
     ) -> None:
         super().__init__(stage=stage)
         self.right_encoder = right_encoder
-        self.left_encoder = left_encoder
+        self.use_shared_encoder = use_shared_encoder
         self.decoder = decoder
-        self.module_list = [right_encoder, left_encoder, decoder]
+        if self.use_shared_encoder:
+            self.module_list = [right_encoder, decoder]
+        else:
+            self.left_encoder = left_encoder
+            self.module_list = [right_encoder, left_encoder, decoder]
         self.output_layer = output_layer
         self.stage = stage
         log_class_usage(__class__)