使用Pytorch构建网络时,损失函数需要迁移到GPU上吗?
·
通常不需要。
简单地说,损失函数接收一个或多个输入tensor,如果输入tensor本身就是在gpu上,则输出tensor自然就在gpu上。
构建网络时,需要主动to或cuda的包括模型输入tensor,以及对应的Module(网络模型)对象。模型输入tensor是数据不用多说,Module对象使用.cuda是为了将模型内部的数据(即每个子叶Module的Parameter或Buffer成员变量)递归地传入gpu。
再进一步,如CrossEntropyCriterion等pytorch自带的loss类(或称loss layer)虽然也是Module的子类,但是这些类的成员变量不含Parameter对象或Buffer对象,所以无需to或cuda处理。Module对象是否持有数据,在设计上也是区分network layer和loss layer的关键。
魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐


所有评论(0)