深度学习实现余弦相似度计算加速-GPU

最新推荐文章于 2024-07-09 02:08:31 发布

绛洞花主敏明

最新推荐文章于 2024-07-09 02:08:31 发布

阅读量3.5k

点赞数 6

分类专栏： Pytorch框架

本文链接：https://blog.csdn.net/qq_39852676/article/details/107342584

版权

Pytorch框架专栏收录该内容

30 篇文章 5 订阅

订阅专栏

一、前言

本文中根据甲方需求实现调用GPU加速计算两千万次的余弦相似度计算，根据需求分析，可以大致分为两个实现目标。

在10s内完成1：2000000次的余弦相似度计算。
在10s内同时完成topk的计算。

二、实现方法

因为在需求中需要完成GPU来加载数据和计算，因此使用普通的程序是无法完成调用GPU的任务，因此，在实现过程中决定采用卷积神经网络的思想来完成数据的加载和计算，同时也能根据需求，分别测试CPU和GPU两种设备下的速度。

初次之外，本文在实现过程中分别采用了两种方式来完成，分别如下：

采用调用torch网络中的torch.nn.function.cosine_similarity来计算余弦相似度
自行构造矩阵同时完成多纬度的计算。
经过测试，第一种方式中使用的for循环严重导致了时间的加长，第二种方式则能分别实现数据的加载和计算，因此，选择第二种方式更加合理。

2.1、方式一：For循环

主要实现分为两步：

构造网络结构，本文中的网络实际仅为余弦相似度的计算。class CosineSimilarity为网络结构。
生成随机特征，并加载到device设备中。
进行计算并循环第二部分。

代码如下：

import torch
import torch.nn.functional as F
import time


class CosineSimilarity(torch.nn.Module):

    def __init__(self, dim=1, eps=1e-8):
        super(CosineSimilarity, self).__init__()
        self.dim = dim
        self.eps = eps

    def forward(self, x1, x2):
        return F.cosine_similarity(x1, x2, self.dim, self.eps)


def test():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = CosineSimilarity().to(device)
    model = torch.nn.DataParallel(model)

    start_time = time.time()
    for i in range(50000):
        x1 = torch.randn(1, 256).to(device)
        x2 = torch.randn(1, 256).to(device)

        distance = model(x1, x2)
        print("index:{}, similar:{}".format(i, distance))
    end_time = time.time()
    print("time is :{}".format(end_time-start_time))

if __name__ == "__main__":
    test()

2.2、方式二：自行构造矩阵

与方式一大致想通过，不同的是第一种方式中每次同时生成两个256纬的数据，加载到设备中之后再进行一次计算，第二种方式是直接全部加载数据到设备中再进行计算。基本过程与第一种方式类似。

构建网络模型，CosineSimilarityTest
加载数据到设备中CPU/GPU
加载模型到设备中model = CosineSimilarityTest().to(device)
进行计算

import torch
import torch.nn.functional as F
import time


class CosineSimilarityTest(torch.nn.Module):
    def __init__(self):
        super(CosineSimilarityTest, self).__init__()

    def forward(self, x1, x2):
        x2 = x2.t()
        x = x1.mm(x2)
    
        x1_frobenius = x1.norm(dim=1).unsqueeze(0).t()
        x2_frobenins = x2.norm(dim=0).unsqueeze(0)
        x_frobenins = x1_frobenius.mm(x2_frobenins)
    
        final = x.mul(1/x_frobenins)
        return final


def main():
    #  加载数据到设备中CP\GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    x1 = torch.randn(5000000, 256).to(device)
    x2 = torch.randn(1, 256).to(device)

    start_time = time.time()
    model = CosineSimilarityTest().to(device)

    # 同时需要多卡计算时需要
    # model = torch.nn.DataParallel(model)

    final_value = model(x1, x2)
    print(final_value.size())

	# 输出排序并输出topk的输出
    value, indec = torch.topk(final_value, 3, dim=0, largest=True, sorted=True)
    print(value)

    end_time = time.time()
    print("消耗时间为：{}".format(end_time - start_time))


if __name__ == "__main__":
    main()