#833·apex

No speedup on RTX card, how apex affects loss function that uses long float?

Author: ulorentzCreated May 12, 2020Updated May 14, 2026

My network is a encoder-decoder type, where the encoder uses resnet blocks (so convolution and batch norm). The decoder uses convtranspose and convolution. I am training on RTX2070super, but the apex training is actually slower than the normal one, what may be the issue since my card has tensorcores? My network requires a custom loss function:

loss = self.loss_func(F.log_softmax(y, 1), yb.long())                                                                                                        

loss1 = self.loss_func(F.log_softmax(y1, 1),                              
                                   F.max_pool2d(yb, kernel_size=2, stride=2,          
                                   padding=0).long())                                                                                                                   
loss2 = self.loss_func(F.log_softmax(y2, 1),                              
                                   F.max_pool2d(yb, kernel_size=4, stride=4,          
                                   padding=0).long())                                                                                                                 

loss3 = self.loss_func(F.log_softmax(y3, 1),                              
                                   F.max_pool2d(yb, kernel_size=8, stride=8,          
                                   padding=0).long())                                                                                                                

loss4 = self.loss_func(F.log_softmax(y4, 1),                              
                                   F.max_pool2d(yb, kernel_size=16, stride=16,        
                                   padding=0).long())                                 
                                                                                 
avg_loss = (loss + (0.9*loss1) + (0.8*loss2) + (0.7*loss3) +              
                    (0.6*loss4))/5   

Where self.loss_func is nn.NLLLoss that appears to require long float as target. My original data target is actually natively float16, may be this conversion a bottleneck? How does apex affect pytorch loss function that use long? Is there any workaround?