通过os.dup sys.stdout.fileno捕获标准输出,判断pytorch算子是否fallback到了cpu

通过os.dup sys.stdout.fileno捕获标准输出,判断pytorch算子是否fallback到了cpu

某种设备在运行pytorch算子时,如果不支持会自动fallback到cpu,输出的tensor.device却不是cpu,我希望能获取到这个状态。本文通过捕获标准输出,根据终端是否输出fallback字符串,判断是否触发了fallback

一.代码

python 复制代码
import threading
import sys
import os

class CheckFallback:
    def __init__(self,enable=True):        
        self.is_fallback=False
        self.enable=enable
        if self.enable:
            self.stdout_fileno_origin = sys.stdout.fileno()
            self.stdout_fileno_dup = os.dup(self.stdout_fileno_origin)
            self.stdout_pipe = os.pipe()
            os.dup2(self.stdout_pipe[1], self.stdout_fileno_origin)
            os.close(self.stdout_pipe[1])
            self.stdout_messages = ''
            self.running=True
            self.task = threading.Thread(target=self.read_pipe)
            self.task.start()

    def read_pipe(self):
        while self.running:
            msg = os.read(self.stdout_pipe[0], 8192)
            if msg:
                self.stdout_messages+=msg.decode('utf-8')
    
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.enable:
            self.running=False
            os.close(self.stdout_fileno_origin)
            self.task.join()
            os.close(self.stdout_pipe[0])
            os.dup2(self.stdout_fileno_dup, self.stdout_fileno_origin)
            os.close(self.stdout_fileno_dup)
            #检查终端是否有fallback信息输出
            if self.stdout_messages.find("fallback")>=0:
                self.is_fallback=True

import torch
A=torch.ones((512,65024),dtype=torch.float16).to("your_device")
with CheckFallback() as f:
    C=torch.ops.aten.gelu.default(A)    
print(f.is_fallback)
print(C.shape,C.device)

with CheckFallback() as f:
    A=torch.ones((1,32),dtype=torch.float16).to("your_device")
    C=torch.ops.aten.pow(A,A)
print(f.is_fallback)
print(C.shape,C.device)
相关推荐
北风toto1 分钟前
python学习DataFrame数据结构
数据结构·python·学习
亿牛云爬虫专家5 分钟前
微服务化采集平台:可扩展性与容错机制
python·微服务·架构·爬虫代理·扩展性·新浪财经·财经新闻
傻啦嘿哟21 分钟前
Python爬虫动态IP代理报错全解析:从问题定位到实战优化
爬虫·python·tcp/ip
mit6.82423 分钟前
[Meetily后端框架] Whisper转录服务器 | 后端服务管理脚本
c++·人工智能·后端·python
zhangfeng113324 分钟前
python 数据分析 单细胞测序数据分析 相关的图表,常见于肿瘤免疫微环境、细胞亚群功能研究 ,各图表类型及逻辑关系如下
开发语言·python·数据分析·医学
Baihai IDP26 分钟前
AI 系统架构的演进:LLM → RAG → AI Workflow → AI Agent
人工智能·ai·系统架构·llm·agent·rag·白海科技
沫儿笙28 分钟前
弧焊机器人气体全方位节能指南
网络·人工智能·机器人
LONGZETECH29 分钟前
【龙泽科技】新能源汽车维护与动力蓄电池检测仿真教学软件【吉利几何G6】
人工智能·科技·汽车·汽车仿真教学软件·汽车教学软件
柠檬豆腐脑43 分钟前
Trae-Agent 内置工具深度解析
python·llm·agent
jndingxin1 小时前
OpenCV 图像哈希类cv::img_hash::AverageHash
人工智能·opencv·哈希算法