美文网首页
JAX 记录

JAX 记录

作者: yxd886 | 来源:发表于2020-10-16 11:15 被阅读0次

    测试官方sample里的resnet50,用的机器是单卡v100.

    batch size 设置为32.
    先测试了一下对update函数默认带了@jit的,也就是开启了XLA Jit优化。
    执行时间大概是0.16s/step

    然后关闭jit发现时间变成了2~3s/step。两者差异巨大。

    于是用nvprof profile了一下。
    使用jit的情况:
    GPU kernel 执行情况:


    image.png

    API call 情况:


    image.png

    关闭jit的情况:

    GPU kernel 情况:


    image.png

    API call 情况:


    image.png

    相关文章

      网友评论

          本文标题:JAX 记录

          本文链接:https://www.haomeiwen.com/subject/qahzpktx.html