Output layer of multiclass 3d image classification

Question

I\'m building a 3d classification NN with TensorFlow. The code was modified from a segmentation code. I\'m now confused about how to build an output layer.

The origin