通过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)
相关推荐
飞哥数智坊10 小时前
GPT-5-Codex 发布,Codex 正在取代 Claude
人工智能·ai编程
倔强青铜三10 小时前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试
虫无涯11 小时前
Dify Agent + AntV 实战:从 0 到 1 打造数据可视化解决方案
人工智能
Dm_dotnet13 小时前
公益站Agent Router注册送200刀额度竟然是真的
人工智能
算家计算13 小时前
7B参数拿下30个世界第一!Hunyuan-MT-7B本地部署教程:腾讯混元开源业界首个翻译集成模型
人工智能·开源
用户25191624271113 小时前
Python之语言特点
python
机器之心13 小时前
LLM开源2.0大洗牌:60个出局,39个上桌,AI Coding疯魔,TensorFlow已死
人工智能·openai
刘立军14 小时前
使用pyHugeGraph查询HugeGraph图数据
python·graphql
Juchecar14 小时前
交叉熵:深度学习中最常用的损失函数
人工智能
林木森ai14 小时前
爆款AI动物运动会视频,用Coze(扣子)一键搞定全流程(附保姆级拆解)
人工智能·aigc