import onnx
from onnx.helper import make_node, make_graph, make_tensor, make_tensor_value_info, make_model
import caffe2.python.onnx.backend as c2

import numpy as np

num = 1 * 256 * 56 * 56
X = np.arange(num).reshape([1, 256, 56, 56]).astype(np.float32)

node_def = make_node(
    'MaxPool',
    inputs=['X'],
    outputs=['Y'],
    kernel_shape=[1, 1],
    strides=[2,2],
    pads=[0,0,0,0])

output = c2.run_node(
        node_def, {"X": X})
print("maxpool:")
print(output.Y.shape)
for i in range(10):
    print(output.Y.flatten()[i])

W = np.full((256,1,1,1),1.0).astype(np.float32)

node_def = make_node(
    'Conv',
    inputs=['X', 'W'],
    outputs=['Y'],
    kernel_shape=[1, 1],
    group=256,
    pads=[0, 0, 0, 0],
    strides=[2,2])

output = c2.run_node(
        node_def, {"X": X, "W": W})
print("conv:")
print(output.Y.shape)
for i in range(10):
    print(output.Y.flatten()[i])