should I use float or classes as output for the final layer in my neural network?

I decided to just implement a float head with an L1Loss in Pytorch and I created a simple but effective synthetic data set to test the implementation. The data set created images into which a number of small squares were randomly drawn. The training label was simply the number of squares divided by 10, a float number with one decimal digit.
The net trained very quickly and to a high degree of precision – the test samples were correct to the one decimal digit.

As to the original question, the runs I made definitely favoured the float over class.

My take on this is that the implementation in classes had a basic imprecision in the assignment of the classes and, perhaps more importantly, the class implementation has no concept of a “metric”. That is, the float implementation, even if it misses the exact match, will try to generate an output label “close” to the input label while the class implementation has no concept of “close”.

One warning with Pytorch. If you are fitting to one float, be sure to encase it in a length 1 vector in the data generator. Pytorch cannot handle a “naked” float number (even though it does become a vector when batches are done). But it doesn’t complain. This cost me a bunch of time.

Leave a Comment