paddle.nn.functional.cross_entropy中的soft_label时间消耗问题

最近在跑模型的时候发现，同样在imagenet1k数据集，256大小的batchsize下，我自己的模型训练一个epoch需要几乎一个小时，而PiT则差不多只需要半个小时。但是将batchsize调整为64之后时间又几乎相同。测试后发现：

单独在本机上测试，经过对比后发现时间开销主要在计算loss时，

使用相同的api
测试代码：

import paddle.nn as nn
import paddle
import paddle.nn.functional as F
from backbone import build_model
import time
from timeit import Timer

paddle.set_device('cpu')

batchsize = 256
t = paddle.randn([batchsize, 3, 224, 224])
label_margin = paddle.rand([batchsize])
label_margin = paddle.cast(label_margin, paddle.int64)

label_soft = paddle.randn([batchsize, 1000])
model = build_model()
out = model(t)

def criterion_margin(out,label_margin):
    result = F.cross_entropy(out,label_margin)

def criterion_soft(out, label_soft):
    result = F.cross_entropy(out, label_soft, soft_label=True)


def main():
    times = 1
    timer1 = Timer('criterion_margin(out, label_margin)',"from __main__ import criterion_margin",globals=globals())
    timer2 = Timer('criterion_soft(out, label_soft)',"from __main__ import criterion_soft",globals=globals())
    print(f"batch size :{batchsize}， repeat times:{times}")

    print(f"margin criterion:{timer1.timeit(times)}")
    print(f"soft criterion:{timer2.timeit(times)}")


if __name__ == "__main__":
    main()

运行结果：

上面为batchsize为64时，下面为batchsize为264时。

可以发现，在batchsize较小时，使用硬标签和软标签花费的时间相近，而当batchsize增大后，软标签所花费的时间小余硬标签所需时间。

当我们反复迭代的时候这个时间将会产生巨大的差距
如：

左侧为普通硬标签所需时间，而右侧为使用了软标签所花费的时间，可以看出硬标签计算开销几乎为软标签的100倍。