I am attempting to build a classifier with a large input space of one hot encoded vectors. The output should be a vector of labels, with 10000 possible labels each. For example, the labels could correspond to something like: (Class 10, Class 992, Class 800, Class 23, Class 30) etc. My first attempt was to one hot encode this output as well, resulting in an output space of size (10000, 5). However, I quickly discovered that to have an output space this large, I would easily exceed the memory on my GPU (16gb).
Question: What is the best way to encode and implement such an output space in PyTorch?
My first thought was to drop the one hot encoding and encode the outputs exactly as I did above, dropping the word class obviously so that the output would be a 5 dimensional vector (10, 992, 800, 23, 30). However, I wasn't sure which loss function to use here: obviously MSE or other such regression losses would not be appropriate. After doing some research, I found someone who recommended using TensorFlow's softmax_cross_entropy_with_logits (sorry I can't seem to find where I originally saw this anymore). Using https://stackoverflow.com/questions/46218566/pytorch-equivalence-for-softmax-cross-entropy-with-logits/59197990#59197990, I was able to implement a criterion in PyTorch after removing the masked parts, that would accept tensors of size (batch_size, 5) for both labels and outputs and return a number.
However, I'm not sure I'm using this correctly. Attempting to pass the same arguments into the TensorFlow version that this is supposed to be equivalent to throws dimension errors. Also, testing this on some fabricated instances don't seem to match up with what my intuition says it should:
label = torch.tensor([[8.0], [8.0]]) # Corresponds to a two instances of class 8
out = torch.tensor([[7.0], [8.0]]) # Corresponds one incorrect prediction
criterion(out, label) # Expecting a nonzero loss because the first prediction was incorrect
# However, returns tensor(0.)
It is not behaving as expected even with a single label, let alone five. I am thinking with five, if I can get the loss function working for one, to do a balanced sum of each individual.