Pytorch深度学习中分段函数计算加速Tips

1.配置环境

使用环境：python3.7
平台：Windows10
工具：pytorch
IDE：PyCharm

2.问题描述

博主在搭建网络过程中，最后一步设计loss时由于要进行分情况讨论，所以采用了for循环加if判断来实现不同情况loss的计算。这样导致模型训练速度大打折扣，为了解决这个问题，思考了一天终于找到解决方法，如果你也遇到同样的问题或许这篇文章可以对你有所启发

将问题简化

我们有一组数据data，根据data计算loss，这个loss可以视为如下的一个分段函数

3.问题解决

首先直接上代码

# author:Hurricane
# date:  2021/11/17
# E-mail:hurri_cane@qq.com


import torch
import random
import time


def calc_1(data):
	out = []
	for k in data:
		if k < 0.5:
			out.append(k ** 2 / 2)
		else:
			out.append(k - 0.5)
	res = sum(out)
	return res


def calc_2(data):
	data_smaller = data < 0.5
	another_data = ~ data_smaller
	res = sum((data ** 2 / 2) * data_smaller) + sum((data - 0.5) * another_data)
	return res


if __name__ == '__main__':
	data = torch.rand(512, 1)
	t1 = time.time()
	res1 = calc_1(data)
	print("calc1 use %f s, the result is %0.2f" % (time.time() - t1, res1.item()))
	t2 = time.time()
	res2 = calc_2(data)
	print("calc2 use %f s, the result is %0.2f" % (time.time() - t2, res2.item()))

	print("All Done!")

计算结果如下：

代码中calc_1采用的就是for循环加if判断来实现，calc_2采用的博主思考到添加mask来实现的方法。可以看到方法2的速度远胜于方法1。这对计算分段函数提升速度上很有帮助。

造成加速的原因博主推测为计算机对于矩阵计算的效率更高，之后遇到分情况讨论的计算方法或许都可以采用这种mask的方法。

4.结束语

如果本文对你有帮助的话还请点赞、收藏一键带走哦，你的支持是我最大的动力！(づ｡◕ᴗᴗ◕｡)づ