通过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)
相关推荐
胡斌附体9 分钟前
linux测试端口是否可被外部访问
linux·运维·服务器·python·测试·端口测试·临时服务器
视觉语言导航27 分钟前
RAL-2025 | 清华大学数字孪生驱动的机器人视觉导航!VR-Robo:面向视觉机器人导航与运动的现实-模拟-现实框架
人工智能·深度学习·机器人·具身智能
**梯度已爆炸**37 分钟前
自然语言处理入门
人工智能·自然语言处理
likeGhee1 小时前
python缓存装饰器实现方案
开发语言·python·缓存
ctrlworks1 小时前
楼宇自控核心功能:实时监控设备运行,快速诊断故障,赋能设备寿命延长
人工智能·ba系统厂商·楼宇自控系统厂家·ibms系统厂家·建筑管理系统厂家·能耗监测系统厂家
项目題供诗1 小时前
黑马python(二十五)
开发语言·python
读书点滴1 小时前
笨方法学python -练习14
java·前端·python
笑衬人心。1 小时前
Ubuntu 22.04 修改默认 Python 版本为 Python3 笔记
笔记·python·ubuntu
BFT白芙堂1 小时前
睿尔曼系列机器人——以创新驱动未来,重塑智能协作新生态(上)
人工智能·机器学习·机器人·协作机器人·复合机器人·睿尔曼机器人
蛋仔聊测试2 小时前
Playwright 中 Page 对象的常用方法详解
python