numpy 实现 pytorch ToTensor 和 Normalize

最新推荐文章于 2026-03-15 15:19:23 发布

原创最新推荐文章于 2026-03-15 15:19:23 发布 · 2.8k 阅读

4 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#pytorch #python #人工智能

收录于

该代码段展示了如何使用numpy实现类似于`transforms.ToTensor`和`transforms.Normalize`的图像数据归一化操作。首先检查均值和标准差是否为numpy数组，如果不是则转换为数组。接着根据一维还是多维调整其形状，然后进行数据除最大值、减均值和除标准差的操作，最后转置数组得到归一化后的图像数据。

def normalize(data, mean, std):
    # transforms.ToTensor, transforms.Normalize的numpy 实现
    if not isinstance(mean, np.ndarray):
        mean = np.array(mean)
    if not isinstance(std, np.ndarray):
        std = np.array(std)
    if mean.ndim == 1:
        mean = np.reshape(mean, (-1, 1, 1))
    if std.ndim == 1:
        std = np.reshape(std, (-1, 1, 1))
    _max = np.max(abs(data))
    _div = np.divide(data, _max)  # i.e. _div = data / _max
    _sub = np.subtract(_div, mean)  # i.e. arrays = _div - mean
    arrays = np.divide(_sub, std)  # i.e. arrays = (_div - mean) / std
    arrays = np.transpose(arrays, (2, 0, 1))
    return arrays