NCHW与NC1HWC0转换

代码：

# -*- coding: utf-8 -*-
import numpy as np

def nchw_to_nc1hwc0(data, batch_size, num_channels, height, width, block_size):
    assert data.shape == (batch_size, num_channels, height, width)
    
    c0 = block_size
    c1 = (num_channels + c0 - 1) // c0
    
    nc1hwc0_data = np.zeros((batch_size, c1, height, width, c0), dtype=data.dtype)
    for b in range(batch_size):
        for i in range(num_channels):
            for j in range(height):
                for k in range(width):
                    nc1hwc0_data[b, i//c0, j, k, i%c0] = data[b, i, j, k]
                        
    return nc1hwc0_data

def nc1hwc0_to_nchw(data, batch_size, num_channels, height, width, block_size):
    assert data.shape == (batch_size, (num_channels + block_size - 1) // block_size, height, width, block_size)
    
    c0 = block_size
    c1 = (num_channels + c0 - 1) // c0
    
    nchw_data = np.zeros((batch_size, num_channels, height, width), dtype=data.dtype)
    for b in range(batch_size):
        for i in range(num_channels):
            for j in range(height):
                for k in range(width):
                    nchw_data[b, i, j, k] = data[b, i//c0, j, k, i%c0]
                        
    return nchw_data

def nchw_to_nc1hwc0_1(data, batch_size, num_channels, height, width, block_size):
    assert data.shape == (batch_size, num_channels, height, width)

    c0 = block_size
    c1 = (num_channels + c0 - 1) // c0
    
    nc1hwc0_data = np.zeros((batch_size, c1, height, width, c0), dtype=data.dtype)
    for b in range(batch_size):
        for i in range(num_channels):
            c1_idx = i // c0
            c0_idx = i % c0
            nc1hwc0_data[b, c1_idx, :, :, c0_idx] = data[b, i, :, :]
            
    return nc1hwc0_data

def nc1hwc0_to_nchw_1(data, batch_size, num_channels, height, width, block_size):
    assert data.shape == (batch_size, (num_channels + block_size - 1) // block_size, height, width, block_size)
    
    c0 = block_size
    c1 = data.shape[1]
    
    nchw_data = np.zeros((batch_size, num_channels, height, width), dtype=data.dtype)
    for b in range(batch_size):
        for i in range(num_channels):
            c1_idx = i // c0
            c0_idx = i % c0
            nchw_data[b, i, :, :] = data[b, c1_idx, :, :, c0_idx]
            
    return nchw_data

batch_size = 6
num_channels = 11  # 可以设置为任何正整数
height = 7
width = 11
block_size = 16

data = np.random.rand(batch_size, num_channels, height, width)

nc1hwc0_data = nchw_to_nc1hwc0(data, batch_size, num_channels, height, width, block_size)
nc1hwc0_data_1 = nchw_to_nc1hwc0_1(data, batch_size, num_channels, height, width, block_size)
print(nc1hwc0_data.shape)  # (2, 3, 8, 8, 16)
print(nc1hwc0_data_1.shape)

nchw_data = nc1hwc0_to_nchw(nc1hwc0_data, batch_size, num_channels, height, width, block_size)
nchw_data_1 = nc1hwc0_to_nchw(nc1hwc0_data, batch_size, num_channels, height, width, block_size)
print(nchw_data.shape)  # (2, 33, 8, 8)
print(nchw_data_1.shape)

assert np.allclose(data, nchw_data)  # 验证转换后得到的数据与原始数据相同
assert np.allclose(data, nchw_data_1)
assert np.allclose(nchw_data, nchw_data_1)