| |
| |
|
|
|
|
|
|
|
|
|
|
| from caffe2.python import schema |
| from caffe2.python.layers.layers import ( |
| ModelLayer, |
| ) |
|
|
|
|
| class PairwiseSimilarity(ModelLayer): |
|
|
| def __init__(self, model, input_record, output_dim, pairwise_similarity_func='dot', |
| name='pairwise_similarity', **kwargs): |
| super(PairwiseSimilarity, self).__init__(model, name, input_record, **kwargs) |
| assert isinstance(input_record, schema.Struct), ( |
| "Incorrect input type. Expected Struct, but received: {0}". |
| format(input_record)) |
| assert ( |
| ('all_embeddings' in input_record) ^ |
| ('x_embeddings' in input_record and 'y_embeddings' in input_record) |
| ), ( |
| "either (all_embeddings) xor (x_embeddings and y_embeddings) " + |
| "should be given." |
| ) |
| self.pairwise_similarity_func = pairwise_similarity_func |
| if 'all_embeddings' in input_record: |
| x_embeddings = input_record['all_embeddings'] |
| y_embeddings = input_record['all_embeddings'] |
| else: |
| x_embeddings = input_record['x_embeddings'] |
| y_embeddings = input_record['y_embeddings'] |
|
|
| assert isinstance(x_embeddings, schema.Scalar), ( |
| "Incorrect input type for x. Expected Scalar, " + |
| "but received: {0}".format(x_embeddings)) |
| assert isinstance(y_embeddings, schema.Scalar), ( |
| "Incorrect input type for y. Expected Scalar, " + |
| "but received: {0}".format(y_embeddings) |
| ) |
|
|
| if 'indices_to_gather' in input_record: |
| indices_to_gather = input_record['indices_to_gather'] |
| assert isinstance(indices_to_gather, schema.Scalar), ( |
| "Incorrect type of indices_to_gather. " |
| "Expected Scalar, but received: {0}".format(indices_to_gather) |
| ) |
| self.indices_to_gather = indices_to_gather |
| else: |
| self.indices_to_gather = None |
|
|
| self.x_embeddings = x_embeddings |
| self.y_embeddings = y_embeddings |
|
|
| dtype = x_embeddings.field_types()[0].base |
|
|
| self.output_schema = schema.Scalar( |
| (dtype, (output_dim,)), |
| self.get_next_blob_reference('output') |
| ) |
|
|
| def add_ops(self, net): |
| if self.pairwise_similarity_func == "cosine_similarity": |
| x_embeddings_norm = net.Normalize(self.x_embeddings(), axis=1) |
| y_embeddings_norm = net.Normalize(self.y_embeddings(), axis=1) |
| Y = net.BatchMatMul( |
| [x_embeddings_norm, y_embeddings_norm], |
| [self.get_next_blob_reference(x_embeddings_norm + '_matmul')], |
| trans_b=1, |
| ) |
| elif self.pairwise_similarity_func == "dot": |
| Y = net.BatchMatMul( |
| [self.x_embeddings(), self.y_embeddings()], |
| [self.get_next_blob_reference(self.x_embeddings() + '_matmul')], |
| trans_b=1, |
| ) |
| else: |
| raise NotImplementedError( |
| "pairwise_similarity_func={} is not valid".format( |
| self.pairwise_similarity_func |
| ) |
| ) |
|
|
| if self.indices_to_gather: |
| flattened = net.Flatten( |
| Y, Y + '_flatten', |
| ) |
| net.BatchGather( |
| [flattened, self.indices_to_gather()], |
| self.output_schema(), |
| ) |
| else: |
| net.Flatten(Y, self.output_schema()) |
|
|
