Deci-AI
/
super-gradients
connected to https://github.com/Deci-AI/super-gradients.git


  
1

	
2

	
3

	
4

	
5

	
6

	
7

	
8

	
9

	
10

	
11

	
12

	
13

	
14

	
15

	
16

	
17

	
18

	
19

	
20

	
21

	
22

	
23

	
24

	
25

	
26

	
27

	
28

	
29

	
30

	
31

	
            from torch import nn
from super_gradients.common.object_names import Losses
from super_gradients.training.losses import (
    LabelSmoothingCrossEntropyLoss,
    ShelfNetOHEMLoss,
    ShelfNetSemanticEncodingLoss,
    RSquaredLoss,
    SSDLoss,
    BCEDiceLoss,
    YoloXDetectionLoss,
    YoloXFastDetectionLoss,
    KDLogitsLoss,
    DiceCEEdgeLoss,
)
from super_gradients.training.losses.stdc_loss import STDCLoss


LOSSES = {
    Losses.CROSS_ENTROPY: LabelSmoothingCrossEntropyLoss,
    Losses.MSE: nn.MSELoss,
    Losses.R_SQUARED_LOSS: RSquaredLoss,
    Losses.SHELFNET_OHEM_LOSS: ShelfNetOHEMLoss,
    Losses.SHELFNET_SE_LOSS: ShelfNetSemanticEncodingLoss,
    Losses.YOLOX_LOSS: YoloXDetectionLoss,
    Losses.YOLOX_FAST_LOSS: YoloXFastDetectionLoss,
    Losses.SSD_LOSS: SSDLoss,
    Losses.STDC_LOSS: STDCLoss,
    Losses.BCE_DICE_LOSS: BCEDiceLoss,
    Losses.KD_LOSS: KDLogitsLoss,
    Losses.DICE_CE_EDGE_LOSS: DiceCEEdgeLoss,
}