现在回顾本文,其实就是一个因为不熟悉库函数而自己去实现的愚蠢故事。
记住,下次自己想实现一些基础功能的时候,一定要先去看看能用上哪些库函数。
目标
在自己写代码的时候,碰到一个蛮有意思的函数。有点像算法题,目标是用torch和numpy的库函数实现。
先介绍一下函数的目标,就是从一个n行m列的数组中,选出每m个中最大的k个,返回一个n*k的数组,每个元素是对应k个元素的下标。返回下标可以乱序。
def select_maxk(arr, pick_num=1)
arr是一个n行m列的二维数组,pick_num就是上面提到的k的值,默认为1。
另外,还有epsilon = 0.1 的概率是随机选择k个元素的下标,来模拟一些特定的场景。(这需要在函数中随机生成一个向量来进行控制。)
比如生成向量tensor([False, True, False, False, False, False]) 就表示,除了第二行是随机返回k个下标,其余几行都必须返回最大元素的下标。
举个例子:n=6,m=4,k=3
原数组:
tensor([[0.1266, 0.1014, 0.1671, 0.6924],
[0.5884, 0.4361, 0.3644, 0.0175],
[0.8387, 0.7729, 0.0654, 0.1425],
[0.5068, 0.7436, 0.3154, 0.3503],
[0.4896, 0.6226, 0.1646, 0.6054],
[0.6791, 0.5005, 0.2655, 0.2428]])
每行最大的3个元素选出来后返回它们在本行的坐标:
tensor([False, True, False, False, False, False])
tensor([[0., 2., 3.],
[3., 1., 2.],
[3., 0., 1.],
[3., 0., 1.],
[0., 1., 3.],
[2., 1., 0.]])
即第一行最大的是第0、2、3个元素,第二行最大的是第0、1、2个元素(但因为对应的向量是True所以随机返回了1、2、3),依次类推。
如果k = 2则返回
tensor([False, False, False, False, False, False])
tensor([[2., 3.],
[1., 0.],
[0., 3.],
[0., 1.],
[3., 1.],
[1., 0.]])
函数实现
import numpy as np
import random
import torch as th
def select_maxk(arr, pick_num=1):
row = arr.shape[0]
line = arr.shape[1]
epsilon = 0.1
random_numbers = th.rand(row)
pick_random = random_numbers < epsilon
# print(pick_random)
result = th.empty(row, pick_num)
for i in range(row):
if pick_random[i]:
randk = th.tensor(random.sample(range(line), pick_num))
result[i] = randk
else:
maxk = th.tensor(np.argpartition(arr[i].numpy().ravel(), -pick_num)[-pick_num:])
result[i] = maxk
return result
这是初稿。我这个实现不是最优的。最起码有这么几点可以改进:
- 随机生成向量这件事可以丢给另外一个函数做
- 不用for循环而采用矩阵的方式
- 混用了numpy和torch,增加数据计算的负担
————————————————
优化改进
其实如果熟悉pytorch的小伙伴,知道topk这个函数的话,直接调库就可以搞定了。
https://pytorch.org/docs/stable/generated/torch.topk.html#torch.topk
torch.topk(input, k, dim=None, largest=True, sorted=True, *, out=None)
>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1., 2., 3., 4., 5.])
>>> torch.topk(x, 3)
torch.return_types.topk(values=tensor([5., 4., 3.]), indices=tensor([4, 3, 2]))
这样的话我上面的算法就是一个伪命题了,不如说,如果一开始就知道这个函数,就没有这篇博客了。
一开始的时候我没有在网上找到类似的接口。结果就让自己费时又费力。
这个故事告诉我们,一定要熟悉库函数,没事多去看看库函数!
优化后函数
import numpy as np
import random
import torch as th
def select_topk(arr, pick_num=1):
row = arr.shape[0]
line = arr.shape[1]
epsilon = 0.1
random_numbers = th.rand(row)
pick_random = random_numbers < epsilon
_, result = th.topk(arr, pick_num, dim = 1)
for i in range(row):
if pick_random[i]:
randk = th.tensor(random.sample(range(line), pick_num))
result[i] = randk
return result
总结
多去看看库函数!
多去看看库函数!
多去看看库函数!