基于picklerpc的pytorch单算子测试[单算子远程测试]

基于picklerpc的pytorch单算子测试[单算子远程测试]

通过主流大模型测试程序-用于导出算子列表 得到了算子类型及参数信息。我希望对比每个算子在不同硬件平台上的性能和误差。如果将所有的结果都存成文件,则占用空间太大。下文演示了如何使用picklerpc 将算子类型及参数传递到远程服务器测试

一.服务端

python 复制代码
from picklerpc import PickleRPCServer
import torch
import numpy as np
import time
import traceback

class TorchOpRunner(PickleRPCServer):
    def __init__(self, addr=('localhost', 8080)):
        super().__init__(addr)

    def run(self,op_type,input_desc):
        input_args=[]
        input_kwargs={}

        for arg in input_desc:
            seed,shape,dtype,device=arg
            torch.random.manual_seed(seed)
            input_args.append(torch.rand(shape,dtype=dtype,device=device))

        op=eval(f"torch.ops.{op_type}")
        warmup_count=1
        test_count=3

        record={}
        record["error"]=0

        try:
            for _ in range(warmup_count):
                output=op(*input_args,**input_kwargs)
            torch.cuda.synchronize()

            t0=time.time()
            for _ in range(test_count):
                output=op(*input_args,**input_kwargs)    
            torch.cuda.synchronize()
            t1=time.time()

            latency=(t1-t0)/test_count

            all=[]
            if isinstance(output,torch.Tensor):
                all.append(output.detach().cpu().float().numpy().reshape(-1))
            elif isinstance(output,list) or isinstance(output,tuple):
                for out in output:
                    if isinstance(out,torch.Tensor):   
                        all.append(out.detach().cpu().float().numpy().reshape(-1))
            else:
                print("error type:",type(output))
                record["error"]=3
            
            if len(all)!=0 and record["error"]==0:
                all=np.concatenate(all,axis=0)
                if all.shape[0]>0:
                    record["data"]=all
                else:
                    record["error"]=5
            else:
                 record["error"]=4
            record["latency"]=latency
            return record
        except:
            traceback.print_exc()
            record["error"]=6
            return record
    
    def raise_error(self):
        """Raise an error"""
        raise NotImplementedError('Not ready')

if __name__ == '__main__':
    srv = TorchOpRunner(addr=('localhost',10001))
    srv.register_function(srv.run)
    srv.serve_forever()

二.客户端

python 复制代码
import torch
import picklerpc

def main():
    op_type="aten.gelu_backward.default"
    seed=0
    shape=(1,512,40,128)
    dtype=torch.float32
    device="cuda:0"
    input_desc=[(seed,shape,dtype,device),(seed,shape,dtype,device)]
    client = picklerpc.PickleRPCClient(('localhost', 10001))
    output=client.run(op_type,input_desc)
    print(output["error"],output["data"].shape)

main()
相关推荐
MicroTech20255 分钟前
微算法科技(NASDAQ:MLGO)开发延迟和隐私感知卷积神经网络分布式推理,助力可靠人工智能系统技术
人工智能·科技·算法
喜欢吃豆11 分钟前
多轮智能对话系统架构方案(可实战):从基础模型到自我优化的对话智能体,数据飞轮的重要性
人工智能·语言模型·自然语言处理·系统架构·大模型·多轮智能对话系统
文火冰糖的硅基工坊26 分钟前
[嵌入式系统-83]:算力芯片的类型与主流架构
人工智能·重构·架构
二王一个今1 小时前
Python打包成exe(windows)或者app(mac)
开发语言·python·macos
一勺菠萝丶1 小时前
Mac 上用 Homebrew 安装 JDK 8(适配 zsh 终端)完整教程
java·python·macos
视觉语言导航2 小时前
ICRA-2025 | 阿德莱德机器人拓扑导航探索!TANGO:具有局部度量控制的拓扑目标可穿越性感知具身导航
人工智能·机器人·具身智能
C嘎嘎嵌入式开发6 小时前
(2)100天python从入门到拿捏
开发语言·python
Stanford_11066 小时前
如何利用Python进行数据分析与可视化的具体操作指南
开发语言·c++·python·微信小程序·微信公众平台·twitter·微信开放平台
西猫雷婶7 小时前
CNN卷积计算
人工智能·神经网络·cnn
white-persist8 小时前
Python实例方法与Python类的构造方法全解析
开发语言·前端·python·原型模式