chaoyuaw commented on a change in pull request #9165: add embedding learning example URL: https://github.com/apache/incubator-mxnet/pull/9165#discussion_r166049971
########## File path: example/gluon/embedding_learning/model.py ########## @@ -0,0 +1,224 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +from mxnet import gluon +from mxnet.gluon import nn, Block, HybridBlock +import numpy as np + +class L2Normalization(HybridBlock): + r"""Applies L2 Normalization to input. + + Parameters + ---------- + mode : str + Mode of normalization. + See :func:`~mxnet.ndarray.L2Normalization` for available choices. + + Inputs: + - **data**: input tensor with arbitrary shape. + + Outputs: + - **out**: output tensor with the same shape as `data`. + """ + def __init__(self, mode, **kwargs): + self._mode = mode + super(L2Normalization, self).__init__(**kwargs) + + def hybrid_forward(self, F, x): + return F.L2Normalization(x, mode=self._mode, name='l2_norm') + + def __repr__(self): + s = '{name}({_mode})' + return s.format(name=self.__class__.__name__, + **self.__dict__) + + +def get_distance(F, x): + """Helper function for margin-based loss. Return a distance matrix given a matrix.""" + n = x.shape[0] + + square = F.sum(x ** 2.0, axis=1, keepdims=True) + distance_square = square + square.transpose() - (2.0 * F.dot(x, x.transpose())) + + # Adding identity to make sqrt work. + return F.sqrt(distance_square + F.array(np.identity(n))) + +class DistanceWeightedSampling(HybridBlock): + r"""Distance weighted sampling. See "sampling matters in deep embedding learning" + paper for details. + + Parameters + ---------- + batch_k : int + Number of images per class. + + Inputs: + - **data**: input tensor with shape (batch_size, embed_dim). + Here we assume the consecutive batch_k examples are of the same class. + For example, if batch_k = 5, the first 5 examples belong to the same class, + 6th-10th examples belong to another class, etc. + + Outputs: + - a_indices: indices of anchors. + - x[a_indices]: sampled anchor embeddings. + - x[p_indices]: sampled positive embeddings. + - x[n_indices]: sampled negative embeddings. + - x: embeddings of the input batch. + """ + def __init__(self, batch_k, cutoff=0.5, nonzero_loss_cutoff=1.4, **kwargs): + self.batch_k = batch_k + self.cutoff = cutoff + + # We sample only from negatives that induce a non-zero loss. + # These are negatives with a distance < nonzero_loss_cutoff. + # With a margin-based loss, nonzero_loss_cutoff == margin + beta. + self.nonzero_loss_cutoff = nonzero_loss_cutoff + super(DistanceWeightedSampling, self).__init__(**kwargs) + + def hybrid_forward(self, F, x): + k = self.batch_k + n, d = x.shape + + distance = get_distance(F, x) + # Cut off to avoid high variance. + distance = F.maximum(distance, self.cutoff) + + # Subtract max(log(distance)) for stability. + log_weights = ((2.0 - float(d)) * F.log(distance) + - (float(d - 3) / 2) * F.log(1.0 - 0.25 * (distance ** 2.0))) + weights = F.exp(log_weights - F.max(log_weights)) + + # Sample only negative examples by setting weights of + # the same-class examples to 0. + mask = np.ones(weights.shape) + for i in range(0, n, k): + mask[i:i+k, i:i+k] = 0 + + weights = weights * F.array(mask) * (distance < self.nonzero_loss_cutoff) + weights = weights / F.sum(weights, axis=1, keepdims=True) + + a_indices = [] + p_indices = [] + n_indices = [] + + np_weights = weights.asnumpy() + for i in range(n): + block_idx = i // k + + try: + n_indices += np.random.choice(n, k-1, p=np_weights[i]).tolist() + except: + n_indices += np.random.choice(n, k-1).tolist() + for j in range(block_idx * k, (block_idx + 1) * k): + if j != i: + a_indices.append(i) + p_indices.append(j) + + return a_indices, x[a_indices], x[p_indices], x[n_indices], x + + def __repr__(self): + s = '{name}({batch_k})' + return s.format(name=self.__class__.__name__, + **self.__dict__) + + +class MarginNet(Block): + r"""Embedding network with distance weighted sampling. + It takes a base CNN and adds an embedding layer and a + sampling layer at the end. + + Parameters + ---------- + base_net : Block + Base network. + emb_dim : int + Dimensionality of the embedding. + batch_k : int + Number of images per class in a batch. Used in sampling. + + Inputs: + - **data**: input tensor with shape (batch_size, channels, width, height). + Here we assume the consecutive batch_k images are of the same class. + For example, if batch_k = 5, the first 5 images belong to the same class, + 6th-10th images belong to another class, etc. + + Outputs: + - The output of DistanceWeightedSampling. + """ + def __init__(self, base_net, emb_dim, batch_k, **kwargs): + super(MarginNet, self).__init__(**kwargs) + with self.name_scope(): + self.base_net = base_net + self.dense = nn.Dense(emb_dim) + self.normalize = L2Normalization(mode='instance') + self.sampled = DistanceWeightedSampling(batch_k=batch_k) + + def forward(self, x): + z = self.base_net(x) + z = self.dense(z) + z = self.normalize(z) + z = self.sampled(z) + return z + + +class MarginLoss(gluon.loss.Loss): + r"""Margin based loss. + + Parameters + ---------- + margin : float + Margin between positive and negative pairs. + nu : float + Regularization parameter for beta. + + Inputs: + - anchors: sampled anchor embeddings. + - positives: sampled positive embeddings. + - negatives: sampled negative embeddings. + - beta_in: class-specific betas. + - a_indices: indices of anchors. Used to get class-specific beta. + + Outputs: + - Loss. + """ + def __init__(self, margin=0.2, nu=0.0, weight=None, batch_axis=0, **kwargs): Review comment: Hi @leezu, thanks for your question! A \nu == 0 doesn't disable learning of \beta^{class}. \nu is the regularization hyperparameter, or you can think of it as setting a prior for beta. In experiments, I found \nu==0 usually works well. Setting --lr-beta=0 will however disable learning of \beta^{class}. I hope that clarifies! ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services
