winograd_cpu.py 文件源码

def calcM(N, Co, U, V):
    GK = U.shape[2]
    Ci = U.shape[3]
    tiles = V.shape[3]
    GN = V.shape[2]

    print('calcM cpu GN', GN, 'N', N)
    U = U.transpose(0,1,2,4,3).reshape(6,6,GK * 32,Ci)[:,:,:Co,:]

    V = V.transpose(
        2,6,0,1,5,3,4).reshape(
        GN * 32, 6, 6, Ci, tiles, tiles)[:N]

    M = np.zeros((N, Co, tiles, tiles, 6, 6), dtype=np.float32)
    for n in range(N):
        for xi in range(6):
            for nu in range(6):
                M[n,:, :, :, xi, nu] = np.tensordot(U[xi,nu], V[n,xi,nu], 1)
    timecheck('calced M')
    return M