def export_to_caffe2(
self,
workspace: core.workspace,
init_net: core.Net,
predict_net: core.Net,
model_out: torch.Tensor,
output_name: str,
) -> List[core.BlobReference]:
"""See `OutputLayerBase.export_to_caffe2()`."""
probability_out = predict_net.Sigmoid(output_name)
return OutputLayerUtils.gen_additional_blobs(
predict_net, probability_out, model_out, output_name, self.target_names
)
def export_to_caffe2(
self,
workspace: core.workspace,
init_net: core.Net,
predict_net: core.Net,
model_out: torch.Tensor,
output_name: str,
) -> List[core.BlobReference]:
"""
Exports the doc classification layer to Caffe2.
See `OutputLayerBase.export_to_caffe2()` for details.
"""
if isinstance(self.loss_fn, BinaryCrossEntropyLoss):
probability_out = predict_net.Sigmoid(output_name)
else:
probability_out = predict_net.Softmax(output_name, axis=model_out.dim() - 1)
return OutputLayerUtils.gen_additional_blobs(
predict_net, probability_out, model_out, output_name, self.target_names
)