julius/prototorch
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Modify stratified_min function

Browse Source
This commit is contained in:
blackfly 2020-04-27 12:48:12 +02:00
parent 532f63b1de
commit d17b9a3346
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
prototorch/functions/competitions.py
View File

@ -7,6 +7,9 @@ import torch
def stratified_min(distances, labels): def stratified_min(distances, labels):
clabels = torch.unique(labels, dim=0) clabels = torch.unique(labels, dim=0)
nclasses = clabels.size()[0] nclasses = clabels.size()[0]
if distances.size()[1] == nclasses:
# skip if only one prototype per class
return distances
batch_size = distances.size()[0] batch_size = distances.size()[0]
winning_distances = torch.zeros(nclasses, batch_size) winning_distances = torch.zeros(nclasses, batch_size)
inf = torch.full_like(distances.T, fill_value=float('inf')) inf = torch.full_like(distances.T, fill_value=float('inf'))