目标:只保留前向功能,尽可能减少代码,以及内存、显存占用。
步骤:
1、 只保留常用的类,包括cpp、cu、hpp文件。
2、 删除类中以Backward 为前缀的函数。
3、 如果使用了cudnn,把CUDNN_STREAMS_PER_GROUP 设置成1。
4、 由于BatchNorm等层的中间结果是独立的,即一个模型中如果有多个BatchNorm层,会产生多份中间变量,可以修改Forward 函数,增加一个保存中间结果的参数,达到复用效果。
5、 去除net类中的大部分成员变量,以及相关的函数。