I have a classifier where the last layer is a torch.nn.Linear (fc) layer with let’s say 1024 output classes.
For genuine inputs which fall into one of the 1024 classes, I want the classifier to output the correct class, i.e. a distribution which can be trained with torch.nn.CrossEntropyLoss such that if I pass the output through torch.softmax() and torch.argmax() I get the correct class.
Now, if I pass bad input, I want the classifier to output a random distribution, for example a gaussian distribution. How do I train that?
So far some candidates are torch.nn.KLDivLoss and torch.nn.GaussianNLLLoss. But I don’t know how to use them in this context. I’ve also thought about calculating the mean, variance, skewness and kurtosis of the output distribution and regressing them to 0, 1, 0 and 3 respectively. But not sure if that’s the right way to do it. There is also the torch.distributions package but not sure what to do with that.
Any ideas?
Thanks