如何更高效地对数组进行独热编码？

拥有约25000个类似下方“原始映射”的数组，需根据数值将每个数组拆分为3个分量（类似灰度图转RGB图），本质是实现独热编码。当前采用嵌套循环的方式实现，效率极低，处理全部数组时会耗尽内存；虽尝试过分批处理，但希望找到更优的数组拆分/独热编码方案，最终用于CNN训练。

原始映射示例 #

# Original Map:
# 0 0 0 1 1 1 0 0 0 0
# 0 0 1 1 1 1 1 0 0 0
# 0 1 1 2 2 2 1 1 0 0
# 1 1 1 2 2 2 1 1 1 0
# 0 1 1 2 2 2 1 1 0 0
# 0 0 1 1 1 1 1 0 0 0
# 0 0 0 1 1 1 0 0 0 0
# 0 0 0 0 0 0 0 0 0 0

# First component[0,;,;,0]   Second component[0,;,;,1] Last component[0,;,;,2]
# (Padding)                   (Good Die)                (Bad Die)

# 1 1 1 0 0 0 1 1 1 1        # 0 0 0 1 1 1 0 0 0 0    # 0 0 0 0 0 0 0 0 0 0
# 1 1 0 0 0 0 0 1 1 1        # 0 0 1 1 1 1 1 0 0 0    # 0 0 0 0 0 0 0 0 0 0
# 1 0 0 0 0 0 0 0 1 1        # 0 1 1 0 0 0 1 1 0 0    # 0 0 0 1 1 1 0 0 0 0
# 0 0 0 0 0 0 0 0 0 1        # 1 1 1 0 0 0 1 1 1 0    # 0 0 0 1 1 1 0 0 0 0
# 1 0 0 0 0 0 0 0 1 1        # 0 1 1 0 0 0 1 1 0 0    # 0 0 0 1 1 1 0 0 0 0
# 1 1 0 0 0 0 0 1 1 1        # 0 0 1 1 1 1 1 0 0 0    # 0 0 0 0 0 0 0 0 0 0
# 1 1 1 0 0 0 1 1 1 1        # 0 0 0 1 1 1 0 0 0 0    # 0 0 0 0 0 0 0 0 0 0
# 1 1 1 1 1 1 1 1 1 1        # 0 0 0 0 0 0 0 0 0 0    # 0 0 0 0 0 0 0 0 0 0

当前低效代码 #

reshaped = np.zeros((1,map_exm.shape[0], map_exm.shape[1], 3))
for h in range(map_exm.shape[0]):
  for w in range(map_exm.shape[1]):
    reshaped[0,h, w, map_exm[h][w]] = 1

1. 用NumPy向量化操作替代嵌套循环 #

NumPy的向量化操作由C底层实现，速度比Python循环快几个数量级，同时内存分配更高效。推荐两种实现方式：

方法1：np.eye直接索引 #

import numpy as np

# 假设map_exm是单个形状为(H, W)的数组
num_classes = 3
# 直接生成独热编码数组，形状为(H, W, 3)
one_hot = np.eye(num_classes)[map_exm]
# 调整维度为(1, H, W, 3)，匹配CNN输入格式
reshaped = one_hot[np.newaxis, ...]

方法2：广播式索引赋值 #

如果担心np.eye创建额外数组占用内存，可直接用广播索引赋值：

import numpy as np

num_classes = 3
reshaped = np.zeros((1, map_exm.shape[0], map_exm.shape[1], num_classes), dtype=np.int8)
# 利用NumPy高级索引直接批量赋值
reshaped[0, np.arange(map_exm.shape[0])[:, None], np.arange(map_exm.shape[1]), map_exm] = 1

注意：用np.int8/np.uint8替代默认的float64，可将内存占用降低75%以上，因为独热编码只有0和1，整数类型完全足够。

2. 批量处理+内存控制 #

如果25000个数组一次性处理内存不足，可结合向量化操作分批处理，比如每次处理100个：

import numpy as np

# 假设所有数据存储在all_maps中，形状为(25000, H, W)
batch_size = 100
num_classes = 3
H, W = all_maps.shape[1], all_maps.shape[2]
processed_batches = []

for i in range(0, len(all_maps), batch_size):
    # 取出当前批次
    batch = all_maps[i:i+batch_size]
    # 批量生成独热编码
    batch_one_hot = np.eye(num_classes)[batch]
    # 或者用广播赋值：
    # batch_one_hot = np.zeros((len(batch), H, W, num_classes), dtype=np.int8)
    # batch_one_hot[np.arange(len(batch))[:, None, None], np.arange(H)[:, None], np.arange(W), batch] = 1
    processed_batches.append(batch_one_hot)

# 合并所有批次（内存允许的情况），否则可直接用批次列表喂给训练流程
processed_data = np.concatenate(processed_batches, axis=0)

3. 用深度学习框架内置函数（适配CNN训练） #

如果用TensorFlow/PyTorch训练，直接用框架自带的独热编码函数，不仅高效，还支持GPU加速，且能和训练流程无缝衔接：

TensorFlow版本 #

import tensorflow as tf

# 若数据是NumPy数组，先转为Tensor
map_tensor = tf.convert_to_tensor(all_maps, dtype=tf.int32)
# 生成独热编码，形状为(25000, H, W, 3)
one_hot = tf.one_hot(map_tensor, depth=3, axis=-1)

# 可直接结合tf.data.Dataset分批加载，无需提前处理全部数据
dataset = tf.data.Dataset.from_tensor_slices(one_hot).batch(32)

PyTorch版本 #

import torch
from torch.utils.data import DataLoader, TensorDataset

# 转为PyTorch LongTensor（one_hot要求输入为整数类型）
map_tensor = torch.tensor(all_maps, dtype=torch.long)
# 生成独热编码，形状为(25000, H, W, 3)
one_hot = torch.nn.functional.one_hot(map_tensor, num_classes=3)
# 调整为PyTorch常用的(批量, 通道, H, W)格式
one_hot = one_hot.permute(0, 3, 1, 2)

# 用DataLoader分批加载
dataset = TensorDataset(one_hot)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

4. 极致内存优化：动态生成独热编码 #

若内存极度紧张，可在训练时动态生成独热编码，无需提前保存全部处理后的数据：

import torch
from torch.utils.data import Dataset, DataLoader

class MapDataset(Dataset):
    def __init__(self, maps):
        self.maps = maps
    
    def __len__(self):
        return len(self.maps)
    
    def __getitem__(self, idx):
        # 取出单张图
        map_data = self.maps[idx]
        # 动态生成独热编码
        one_hot = torch.nn.functional.one_hot(torch.tensor(map_data, dtype=torch.long), num_classes=3)
        # 转为PyTorch训练所需的(通道, H, W)格式
        return one_hot.permute(2, 0, 1)

# 加载数据，batch_size按需设置
dataset = MapDataset(all_maps)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

这种方式每次仅加载一个批次的数据并生成独热编码，内存占用极小。

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