pytorch_cnn_trainer
pytorch_cnn_trainer copied to clipboard
Add Support for multi label classification
🚀 Feature
Should take some time.
We can add an argument while creating models, num_labels: int = 1
. If the user specifies that, then simply we can create an extra dense layer on with the top and return it in the model.
Also, we need to slightly alter the engine to support this. We will get two outs, o1, o2. We would need to compute metrics for both and provide other functions too.
Either we can use the same train_step and val_step or use different too.