from mxnet import gluon,init
from mxnet import autograd,nd
from mxnet.gluon import nn,loss as gloss
from mxnet.gluon import data as gdata


# 二維卷積層
def corr2d(X,K):
    h, w = K.shape
    Y = nd.zeros((X.shape[0] - h + 1,X.shape[1] - w + 1))
    for i in range(Y.shape[0]):
        
 
for j in range(Y.shape[1]):
            Y[i,j] = (X[i: i+h,j:j+w]*K).sum()
    return Y

# 多通道輸入
def corr2d_multi_in(X,K):
    return nd.add_n(*[corr2d(x,k) for x,k in zip(X,K)])


def corr2d_multi_in_out(X,K):
    return nd.stack(*[corr2d_multi_in(X,k) for k in K])

# 1 * 1 卷積層

def corr2d_multi_in_out_1x1(X,K):
    c_i, h, w 
 
= X.shape
    c_o = K.shape[0]
    X = X.reshape((c_i,h*w))
    K = K.reshape((c_o,c_i))
    Y = nd.dot(K,X)
    return Y.reshape((c_o,h,w))

X = nd.random.uniform(shape=(3,3,3))
K = nd.random.uniform(shape=(2,3,1,1))

Y1 = corr2d_multi_in_out_1x1(X, K)
Y2 = corr2d_multi_in_out(X,K)

print
 
((Y1-Y2).norm().asscalar() < 1e-6)