is it possible to change the loss function in classification?

[fyears]

Obviously, in

model2vec/model2vec/train/classifier.py

Line 407 in e0bb55f

nn.CrossEntropyLoss(weight=class_weight)

the loss function is cross entropy.

But I would like to use another loss function (focal loss for imbalanced data)

Is it possible?

[Pringled]

Hi @fyears ,
Yes, I believe you can easily drop in a different loss function if you want (though it's not something that's directly supported as a parameter). The training setup is pretty modular (see https://github.com/MinishLab/model2vec/tree/main/model2vec/train#bring-your-own-architecture for example). I think you can just swap the loss for your desired loss function and it should work, though you might need to update the inputs/output a little bit depending on the function you use.

If this works well for your usecase I'm open to adding this as an official feature as well.