I am working on very similar problem currently. You can use triplet margin loss: http://pytorch.org/docs/nn.html#tripletmarginloss 66. It works out of the box pretty well for me with thousands of classes and small number of examples per class. One network is used to compute feature vectors for anchor, positive and negative example. Then you split the output into 3 matrices and pass them into loss function.