通过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)
相关推荐
蒋星熠7 分钟前
WebSocket网络编程深度实践:从协议原理到生产级应用
网络·数据库·redis·python·websocket·网络协议·微服务
带娃的IT创业者10 分钟前
实战:用 Python 搭建 MCP 服务 —— 模型上下文协议(Model Context Protocol)应用指南
开发语言·python·mcp
Hello123网站12 分钟前
探迹SalesGPT
人工智能·ai工具
摘星星的屋顶14 分钟前
论文阅读记录之《VelocityGPT 》
论文阅读·人工智能·深度学习·学习
万粉变现经纪人14 分钟前
如何解决pip安装报错ModuleNotFoundError: No module named ‘python-dateutil’问题
开发语言·ide·python·pycharm·pandas·pip·httpx
格林威28 分钟前
工业相机如何通过光度立体成像技术实现高效精准的2.5D缺陷检测
人工智能·深度学习·数码相机·yolo·计算机视觉
跟橙姐学代码28 分钟前
Python 类的正确打开方式:从新手到进阶的第一步
前端·python·ipython
c8i29 分钟前
关于python中的变量中使用的下划线_总结
python
MarkHD34 分钟前
大语言模型入门指南:从原理到实践应用
人工智能·语言模型·自然语言处理
A尘埃35 分钟前
NLP(自然语言处理, Natural Language Processing)
人工智能·自然语言处理·nlp