如何无需转换为NumPy直接将大型CuPy数组保存为JPEG？

以下是几个无需转换为NumPy数组、全程（或近乎全程）在GPU上完成的高效方案，完全规避CPU拷贝的性能瓶颈：

1. 利用CuPy扩展库的原生图像保存 #

cupyx.scipy.misc.imsave是CuPy对SciPy工具的GPU实现，直接支持CuPy数组写入图像，无需回传CPU：

import cupy as cp
from cupyx.scipy.misc import imsave

# 生成大型RGB格式CuPy数组（需为uint8类型，范围0-255）
large_img = cp.random.randint(0, 256, size=(4096, 4096, 3), dtype=cp.uint8)

# 直接保存为JPEG
imsave("output.jpg", large_img)

注意： 如果数组是float类型，先在GPU上完成归一化和类型转换：

float_img = cp.random.rand(4096, 4096, 3)  # 0-1范围的float数组
uint8_img = (float_img * 255).clip(0, 255).astype(cp.uint8)
imsave("output_float.jpg", uint8_img)

2. 通过PyTorch的GPU图像IO实现零拷贝保存 #

如果项目已引入PyTorch，可以利用CuPy与PyTorch CUDA张量的零拷贝互操作性，用torchvision.io.write_jpeg直接保存：

import cupy as cp
import torch
import torchvision.io

# CuPy数组转PyTorch CUDA张量（无CPU拷贝）
large_img_cp = cp.random.randint(0, 256, size=(4096, 4096, 3), dtype=cp.uint8)
# 转换为PyTorch要求的CHW格式（通道在前）
large_img_torch = torch.as_tensor(large_img_cp.transpose(2,0,1), device='cuda')

# 直接保存JPEG
torchvision.io.write_jpeg(large_img_torch, "output_torch.jpg", quality=95)

3. 调用NVJPEG底层库（性能最优方案） #

NVJPEG是NVIDIA官方的GPU加速JPEG编解码库，CuPy提供了直接调用的接口，这是目前性能最高的方案，仅在最后将压缩后的JPEG字节流拷贝到CPU写入文件：

import cupy as cp
from cupy.cuda import nvjpeg

# 初始化NVJPEG编码器
encoder = nvjpeg.NVJPEG()

# 准备HWC格式的uint8 CuPy数组
large_img = cp.random.randint(0, 256, size=(8192, 8192, 3), dtype=cp.uint8)

# GPU端完成JPEG编码
jpeg_bytes = encoder.encode(large_img, format=nvjpeg.NVJPEG_INPUT_RGB)

# 仅拷贝压缩后的字节流到CPU并写入文件
with open("output_nvjpeg.jpg", "wb") as f:
    f.write(cp.asnumpy(jpeg_bytes))

优势： 压缩过程完全在GPU上完成，压缩后的字节流数据量仅为原数组的1/20~1/50，CPU拷贝的开销可以忽略不计。

关键注意事项 #

• 所有方案都要求数组为uint8类型且数值范围在0-255，JPEG不支持其他位深的图像，类型转换务必在GPU上完成（用CuPy的astype和clip）。
• 绝对避免使用imageio、OpenCV的默认接口，这些库强制要求NumPy数组，会触发全量CPU拷贝，完全抵消GPU加速的意义。
• 读取JPEG回CuPy数组时，同样使用cupyx.scipy.misc.imread或nvjpeg.decode，直接得到GPU数组。

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