通过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)
相关推荐
众趣科技32 分钟前
数字孪生重构智慧园区:众趣科技何以成为 VR 园区领域标杆
人工智能·3d·智慧城市·空间计算
心勤则明1 小时前
Spring AI 会话记忆实战:从内存存储到 MySQL + Redis 双层缓存架构
人工智能·spring·缓存
木子杳衫1 小时前
【软件开发】管理类系统
python·web开发
ARM+FPGA+AI工业主板定制专家3 小时前
基于GPS/PTP/gPTP的自动驾驶数据同步授时方案
人工智能·机器学习·自动驾驶
长鸳词羡4 小时前
wordpiece、unigram、sentencepiece基本原理
人工智能
ㄣ知冷煖★4 小时前
【GPT5系列】ChatGPT5 提示词工程指南
人工智能
科士威传动4 小时前
丝杆支撑座在印刷设备如何精准运行?
人工智能·科技·自动化·制造
程序员小远5 小时前
银行测试:第三方支付平台业务流,功能/性能/安全测试方法
自动化测试·软件测试·python·功能测试·测试工具·性能测试·安全性测试
taxunjishu5 小时前
DeviceNet 转 Modbus TCP 协议转换在 S7-1200 PLC化工反应釜中的应用
运维·人工智能·物联网·自动化·区块链
kalvin_y_liu6 小时前
智能体框架大PK!谷歌ADK VS 微软Semantic Kernel
人工智能·microsoft·谷歌·智能体