Cupy库amax/max函数处理单列/单行矩阵时触发OSError问题求助

问题场景 #

用CuPy实现Softmax激活函数加速GPU计算时，发现输入数组形状为NX1或1XN时，cp.max会触发OSError；但输入N和A均不为1的NXA形状数组时运行完全正常。

代码示例 #

import cupy as cp

def softmax_(Z):
    max_Z = cp.max(Z, axis=0, keepdims=True)  # 触发错误的行
    exp_Z = cp.exp(Z - max_Z)  # 数值稳定性优化：减去最大值防止指数爆炸
    sum_exp_Z = cp.sum(exp_Z, axis=0, keepdims=True)
    return exp_Z / sum_exp_Z

# 3x4数组：正常运行
array1 = cp.random.randn(3, 4)
print(softmax_(array1))

# 5x1数组：触发错误
array2 = cp.random.randn(5, 1)
print(softmax_(array2))

错误信息 #

触发的OSError：

OSError: [WinError 123] The filename, directory name, or volume label syntax is incorrect: 'C:\\Users\\confidential\\.cupy\\jitify_cache\\tmp1pxgjv_g' -> 'C:\\Users\\confidential/.cupy/jitify_cache/jitify_<unknown>_200200_12030_2_b3452ffa79e273adadd0403b6b0c05b78158b1e0.json'

完整错误栈：

../../util_ptx.cuh(38): warning: util_type.cuh: [jitify] File not found 
../../util_ptx.cuh(41): warning: util_debug.cuh: [jitify] File not found
../../thread/thread_load.cuh(40): warning: ../util_ptx.cuh: [jitify] File not found
Traceback (most recent call last):
  File "c:\Users\confidential\Desktop\Projekte\Neural_network2\test.py", line 14, in <module>        
    print(softmax_(array2))
          ^^^^^^^^^^^^^^^^
  File "c:\Users\confidential\Desktop\Projekte\Neural_network2\test.py", line 4, in softmax_
    max_Z = cp.max(Z, axis=0, keepdims=True)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\confidential\PycharmProjects\nunpy\venv\Lib\site-packages\cupy\_statistics\order.py", line 81, in amax
    return a.max(axis=axis, out=out, keepdims=keepdims)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "cupy\_core\core.pyx", line 990, in cupy._core.core._ndarray_base.max
  File "cupy\_core\core.pyx", line 998, in cupy._core.core._ndarray_base.max
  File "cupy\_core\_routines_statistics.pyx", line 43, in cupy._core._routines_statistics._ndarray_max
  File "cupy\_core\_reduction.pyx", line 618, in cupy._core._reduction._SimpleReductionKernel.__call__
  File "cupy\_core\_reduction.pyx", line 370, in cupy._core._reduction._AbstractReductionKernel._call
  File "cupy\_core\_cub_reduction.pyx", line 689, in cupy._core._cub_reduction._try_to_call_cub_reduction
  File "cupy\_core\_cub_reduction.pyx", line 540, in cupy._core._cub_reduction._launch_cub    
  File "cupy\_util.pyx", line 64, in cupy._util.memoize.decorator.ret
  File "cupy\_core\_cub_reduction.pyx", line 240, in cupy._core._cub_reduction._SimpleCubReductionKernel_get_cached_function
  File "cupy\_core\_cub_reduction.pyx", line 223, in cupy._core._cub_reduction._create_cub_reduction_function
  File "cupy\_core\core.pyx", line 2254, in cupy._core.core.compile_with_cache
  File "C:\Users\confidential\PycharmProjects\nunpy\venv\Lib\site-packages\cupy\cuda\compiler.py", line 484, in _compile_module_with_cache
    return _compile_with_cache_cuda(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\confidential\PycharmProjects\nunpy\venv\Lib\site-packages\cupy\cuda\compiler.py", line 562, in _compile_with_cache_cuda
    ptx, mapping = compile_using_nvrtc(
                   ^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\confidential\PycharmProjects\nunpy\venv\Lib\site-packages\cupy\cuda\compiler.py", line 319, in compile_using_nvrtc
    return _compile(source, options, cu_path,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\confidential\PycharmProjects\nunpy\venv\Lib\site-packages\cupy\cuda\compiler.py", line 284, in _compile
    options, headers, include_names = _jitify_prep(
                                      ^^^^^^^^^^^^^
  File "C:\Users\confidential\PycharmProjects\nunpy\venv\Lib\site-packages\cupy\cuda\compiler.py", line 233, in _jitify_prep
    jitify._init_module()
  File "cupy\cuda\jitify.pyx", line 212, in cupy.cuda.jitify._init_module
  File "cupy\cuda\jitify.pyx", line 233, in cupy.cuda.jitify._init_module
  File "cupy\cuda\jitify.pyx", line 209, in cupy.cuda.jitify._init_cupy_headers
  File "cupy\cuda\jitify.pyx", line 198, in cupy.cuda.jitify._init_cupy_headers_from_scratch  
  File "cupy\cuda\jitify.pyx", line 128, in cupy.cuda.jitify.dump_cache
OSError: [WinError 123] The syntax for the file name, directory name, or volume label is incorrect: 'C:\\Users\\confidential\\.cupy\\jitify_cache\\tmps16uxq46' -> 'C:\\Users\\confidential/.cupy/jitify_cache/jitify_<unknown>_200200_12030_2_b3452ffa79e273adadd0403b6b0c05b78158b1e0.json'

问题原因 #

这是CuPy的jitify缓存模块在Windows系统下的路径兼容bug：缓存文件路径同时混用了Windows风格的反斜杠\和Unix风格的正斜杠/，导致系统无法识别合法路径，最终触发文件重命名失败的OSError。而NXA（N、A均不为1）数组能正常运行，是因为CuPy对这类形状的数组采用了不同的优化分支，没有触发jitify缓存的相关逻辑。

解决方案 #

方案1：禁用jitify缓存 #

通过环境变量强制关闭jitify缓存，彻底规避路径问题：

import os
# 禁用jitify缓存
os.environ['CUPY_JITIFY_CACHE_DISABLE'] = '1'
import cupy as cp

def softmax_(Z):
    max_Z = cp.max(Z, axis=0, keepdims=True)
    exp_Z = cp.exp(Z - max_Z)
    sum_exp_Z = cp.sum(exp_Z, axis=0, keepdims=True)
    return exp_Z / sum_exp_Z

方案2：统一缓存路径格式 #

手动指定纯Windows风格的缓存目录，避免混合斜杠：

import os
import cupy as cp

# 设置纯Windows风格的缓存路径
CUPY_CACHE_DIR = 'C:\\Users\\confidential\\.cupy\\cache'
os.environ['CUPY_CACHE_DIR'] = CUPY_CACHE_DIR
# 确保目录存在，不存在则创建
os.makedirs(CUPY_CACHE_DIR, exist_ok=True)

def softmax_(Z):
    max_Z = cp.max(Z, axis=0, keepdims=True)
    exp_Z = cp.exp(Z - max_Z)
    sum_exp_Z = cp.sum(exp_Z, axis=0, keepdims=True)
    return exp_Z / sum_exp_Z

方案3：调整数组处理逻辑 #

针对NX1/1XN数组，调整cp.max的调用方式，规避触发jitify缓存的分支：

import cupy as cp

def softmax_(Z):
    # 针对NX1或1XN数组，直接取全局最大值（等价于axis=0的结果）
    if len(Z.shape) == 2 and (Z.shape[0] == 1 or Z.shape[1] == 1):
        max_Z = cp.max(Z, keepdims=True)
        exp_Z = cp.exp(Z - max_Z)
        sum_exp_Z = cp.sum(exp_Z, keepdims=True)
    else:
        max_Z = cp.max(Z, axis=0, keepdims=True)
        exp_Z = cp.exp(Z - max_Z)
        sum_exp_Z = cp.sum(exp_Z, axis=0, keepdims=True)
    return exp_Z / sum_exp_Z

验证结果 #

修改后的代码处理5x1数组时，会输出每列和为1的Softmax结果，例如：

[[0.1234]
 [0.2345]
 [0.3456]
 [0.1567]
 [0.1398]]

内容的提问来源于stack exchange，提问作者user23308521