def template(x_shape, axes, keepdims=None, description: str = ""): vx = np.random.rand(*x_shape) vy = np.sum(vx, axis=tuple(axes), keepdims=True if keepdims is None else keepdims) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) kwargs = {"axes": axes} if keepdims is not None: kwargs["keepdims"] = keepdims operator = make_node("ReduceSum", ["x"], ["y"], **kwargs) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) assert tuple(vy.shape) == tuple( graph.outputs[0].shape ), f"vy: {vy.shape}, graph.outputs[0]: {graph.outputs[0].shape}" generate_kernel_test_case( description=f"[ONNX] ReduceSum {description}", graph=graph, backend=["webgpu", "webgl", "webassembly"], inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(x_shape, axis, description: str = ""): np_axis = 1 if axis is None else axis vx = np.random.rand(*x_shape) - 0.5 new_shape = [mul(vx.shape[:np_axis]), mul(vx.shape[np_axis:])] vy = vx.reshape(new_shape) vy = np.exp(vy) / np.sum(np.exp(vy), axis=1, keepdims=True) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) kwargs = {} if axis is not None: kwargs["axis"] = axis operator = make_node("Softmax", ["x"], ["y"], **kwargs) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) assert tuple(vy.shape) == tuple( graph.outputs[0].shape ), f"vy: {vy.shape}, graph.outputs[0]: {graph.outputs[0].shape}" generate_kernel_test_case( description=f"[ONNX] Softmax {description}", graph=graph, backend=["webgpu", "webgl", "webassembly"], inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(x_shape, description: str = ""): vx = np.random.rand(*x_shape) + 1.0 vy = np.sqrt(vx) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) operator = make_node("Sqrt", ["x"], ["y"]) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case(description=f"[ONNX] Sqrt {description}", graph=graph, inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy})
def template(M=5, K=6, N=7, c_shape=None, transA=False, transB=False, broadcast=False, alpha=1.0, beta=0.0, description: str = ""): if c_shape is None: c_shape = [M, N] va = np.random.rand(*((K, M) if transA else (M, K))) vb = np.random.rand(*((N, K) if transB else (K, N))) vc = np.random.rand(*c_shape) vd = (va.T if transA else va) @ (vb.T if transB else vb) * alpha + vc * beta kwargs = { "broadcast": broadcast, "transA": transA, "transB": transB, "alpha": alpha, "beta": beta } a = make_tensor_value_info("a", va.shape) b = make_tensor_value_info("b", vb.shape) c = make_tensor_value_info("c", vc.shape) d = make_tensor_value_info("d", vd.shape) operator = make_node("Gemm", ["a", "b", "c"], ["d"], **kwargs) model = make_model([operator], [a, b, c], [d]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] Gemm {description}", graph=graph, inputs={ graph.inputs[0]: va, graph.inputs[1]: vb, graph.inputs[2]: vc }, expected={graph.outputs[0]: vd}, )
def template(N=2, H=5, W=5, C=7, epsilon: float = 1e-5, description: str = ""): # spatial = 0 is not supported x_shape = [N, C, H, W] def expand(v): return v[np.newaxis, :, np.newaxis, np.newaxis] # computing test case with float64 makes difference def randfloat(*shape): return np.random.rand(*shape).astype(np.float32) vx = randfloat(*x_shape) vscale = randfloat(C) vB = randfloat(C) vmean = randfloat(C) vvar = randfloat(C) vy = (vx - expand(vmean)) / np.sqrt(expand(vvar) + epsilon) * expand(vscale) + expand(vB) x = make_tensor_value_info("x", x_shape) scale = make_tensor_value_info("scale", [C]) B = make_tensor_value_info("B", [C]) mean = make_tensor_value_info("mean", [C]) var = make_tensor_value_info("var", [C]) y = make_tensor_value_info("y", vy.shape) kwargs = {"epsilon": epsilon} operator = make_node("BatchNormalization", ["x", "scale", "B", "mean", "var"], ["y"], **kwargs) model = make_model([operator], [x, scale, B, mean, var], [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] BatchNormalization {description}", graph=graph, inputs={ graph.inputs[0]: vx, graph.inputs[1]: vscale, graph.inputs[2]: vB, graph.inputs[3]: vmean, graph.inputs[4]: vvar }, expected={graph.outputs[0]: vy}, )
def template(x_shape, description: str = ""): vx = np.random.rand(*x_shape) vy = np.mean(vx, axis=tuple(range(2, vx.ndim)), keepdims=False) x = make_tensor_value_info("x", x_shape) y = make_tensor_value_info("y", vy.shape) operator = make_node("GlobalAveragePool", ["x"], ["y"]) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] GlobalAveragePool {description}", graph=graph, backend=["webgpu", "webgl", "webassembly"], inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(x_shape, alpha, description: str = ""): vx = np.random.rand(*x_shape) - 0.5 vy = np.maximum(vx, vx * alpha) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) operator = make_node("LeakyRelu", ["x"], ["y"], alpha=alpha) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] LeakyRelu {description}", graph=graph, inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(N=2, H=5, W=5, C=7, KH=3, KW=3, SH=1, SW=1, PH=1, PW=1, DH=1, DW=1, description: str = ""): if DH != 1 or DW != 1: raise NotImplementedError x_shape = [N, C, H, W] vx = np.random.rand(*x_shape) vy = chainer.functions.max_pooling_2d(vx, ksize=[KH, KW], stride=[SH, SW], pad=[PH, PW]).data x = make_tensor_value_info("x", x_shape) y = make_tensor_value_info("y", vy.shape) kwargs = { "kernel_shape": [KH, KW], "strides": [SH, SW], "dilations": [DH, DW], "pads": [PH, PH, PW, PW] } operator = make_node("MaxPool", ["x"], ["y"], **kwargs) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] MaxPool {description}", graph=graph, inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(xs_shape, axis, description: str = ""): vxs = [np.random.rand(*x_shape) for x_shape in xs_shape] vy = np.concatenate(vxs, axis) xs = [make_tensor_value_info(f"x{i}", vx.shape) for i, vx in enumerate(vxs)] y = make_tensor_value_info("y", vy.shape) operator = make_node("Concat", [x.name for x in xs], ["y"], axis=axis) model = make_model([operator], xs, [y]) graph = ONNXConverter().convert(model) assert tuple(vy.shape) == tuple(graph.outputs[0].shape), f"vy: {vy.shape}, graph.outputs[0]: {graph.outputs[0].shape}" generate_kernel_test_case( description=f"[ONNX] Concat {description}", graph=graph, inputs={graph.inputs[i]: vx for i, vx in enumerate(vxs)}, expected={graph.outputs[0]: vy}, )
def template(x_shape, axis, description: str = ""): vx = np.random.rand(*x_shape) - 0.5 vy = np.exp(vx) / np.sum(np.exp(vx), axis=axis, keepdims=True) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) operator = make_node("Softmax", ["x"], ["y"], axis=axis) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] Softmax {description}", graph=graph, backend=["webgpu", "webgl", "webassembly"], inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(x_shape, axes, description: str = ""): vx = np.random.rand(*x_shape) vy = np.squeeze(vx, tuple(axes)) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) operator = make_node("Squeeze", ["x"], ["y"], axes=axes) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) assert list(vy.shape) == list(graph.outputs[0].shape) generate_kernel_test_case( description=f"[ONNX] Squeeze {description}", graph=graph, inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(x_shape, y_shape, description: str = ""): vx = np.random.rand(*x_shape) vy = np.reshape(vx, [x if y == 0 else y for x, y in zip(x_shape, y_shape)]) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) operator = make_node("Reshape", ["x"], ["y"], shape=y_shape) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) assert list(vy.shape) == list(graph.outputs[0].shape) generate_kernel_test_case( description=f"[ONNX] Reshape {description}", graph=graph, inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(x_shape, split, axis, description: str = ""): vx = np.random.rand(*x_shape) sections = np.cumsum(split).tolist()[:-1] vys = np.split(vx, sections, axis=axis) x = make_tensor_value_info("x", vx.shape) ys = [make_tensor_value_info(f"y{i}", vy.shape) for i, vy in enumerate(vys)] operator = make_node("Split", ["x"], [y.name for y in ys], axis=axis, split=split) model = make_model([operator], [x], ys) graph = ONNXConverter().convert(model) for i, vy in enumerate(vys): assert tuple(vy.shape) == tuple(graph.outputs[i].shape), f"vys[{i}]: {vy.shape}, graph.outputs[{i}]: {graph.outputs[i].shape}" generate_kernel_test_case( description=f"[ONNX] Split {description}", graph=graph, inputs={graph.inputs[0]: vx}, expected={graph.outputs[i]: vy for i, vy in enumerate(vys)}, )
def template(x_shape, axis, kwargs, description: str = ""): vx = np.random.rand(*x_shape) y_shape = [np.product(vx.shape[:axis]), np.product(vx.shape[axis:])] vy = vx.reshape(y_shape) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) operator = make_node("Flatten", ["x"], ["y"], **kwargs) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) assert tuple(vy.shape) == tuple( graph.outputs[0].shape ), f"vy: {vy.shape}, graph.outputs[0]: {graph.outputs[0].shape}" generate_kernel_test_case( description=f"[ONNX] Flatten {description}", graph=graph, inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(x0_shape, x1_shape, broadcast=0, axis=None, description: str = ""): vx0 = np.random.rand(*x0_shape) vx1 = np.random.rand(*x1_shape) if axis is not None: # onnx-style broadcast vx1 = vx1[(None, ) * axis + (..., ) + (None, ) * (vx0.ndim - vx1.ndim - axis)] vy = vx0 + vx1 x0 = make_tensor_value_info("x0", x0_shape) x1 = make_tensor_value_info("x1", x1_shape) y = make_tensor_value_info("y", vy.shape) kwargs = {"broadcast": broadcast} if axis is not None: kwargs["axis"] = axis operator = make_node("Add", ["x0", "x1"], ["y"], **kwargs) model = make_model([operator], [x0, x1], [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] Add {description}", graph=graph, inputs={ graph.inputs[0]: vx0, graph.inputs[1]: vx1 }, expected={graph.outputs[0]: vy}, )
def template(x_shape, pads, mode, value=None, description: str = ""): vx = np.random.rand(*x_shape) np_pads = [[b, e] for b, e in zip(pads[:vx.ndim], pads[vx.ndim:])] if mode == b"constant": vy = np.pad(vx, np_pads, mode="constant", constant_values=value) elif mode == b"reflect": vy = np.pad(vx, np_pads, mode="reflect") elif mode == b"edge": vy = np.pad(vx, np_pads, mode="symmetric") else: raise ValueError(mode) x = make_tensor_value_info("x", vx.shape) y = make_tensor_value_info("y", vy.shape) kwargs = {"pads": pads, "mode": mode} if value is not None: kwargs["value"] = value operator = make_node("Pad", ["x"], ["y"], **kwargs) model = make_model([operator], [x], [y]) graph = ONNXConverter().convert(model) assert list(vy.shape) == list(graph.outputs[0].shape) generate_kernel_test_case( description=f"[ONNX] Pad {description}", graph=graph, backend=["webgpu", "webassembly", "webgl"], inputs={graph.inputs[0]: vx}, expected={graph.outputs[0]: vy}, )
def template(x_shape, num_input=2, description: str = ""): vxs = [np.random.rand(*x_shape) for _ in range(num_input)] vy = np.zeros(x_shape) for vx in vxs: vy += vx xs = [ make_tensor_value_info(f"x{i}", vxs[i].shape) for i in range(num_input) ] y = make_tensor_value_info("y", vy.shape) operator = make_node("Sum", [x.name for x in xs], ["y"]) model = make_model([operator], xs, [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] Sum {description}", graph=graph, inputs={v: x for v, x in zip(graph.inputs, vxs)}, expected={graph.outputs[0]: vy}, )
def template(x0_shape, x1_shape, description: str = ""): vx0 = np.random.rand(*x0_shape) vx1 = np.random.rand(*x1_shape) vy = vx0**vx1 x0 = make_tensor_value_info("x0", vx0.shape) x1 = make_tensor_value_info("x1", vx1.shape) y = make_tensor_value_info("y", vy.shape) operator = make_node("Pow", ["x0", "x1"], ["y"]) model = make_model([operator], [x0, x1], [y]) graph = ONNXConverter().convert(model) generate_kernel_test_case( description=f"[ONNX] Pow {description}", graph=graph, inputs={ graph.inputs[0]: vx0, graph.inputs[1]: vx1 }, expected={graph.outputs[0]: vy}, )
def main(): # Training settings parser = argparse.ArgumentParser(description='PyTorch MNIST Example') parser.add_argument( '--out', '-o', default='output_pytorch', help='Directory to output the graph descriptor and sample test data') parser.add_argument("--backend", default="webgpu,webgl,webassembly,fallback") args = parser.parse_args() training_dir = os.path.join(args.out, "pytorch_model") os.makedirs(training_dir, exist_ok=True) model_path = os.path.join(training_dir, "model.proto") if not os.path.exists(model_path): # model training part (as usual) torch.manual_seed(1) device = torch.device("cpu") train_loader = torch.utils.data.DataLoader( datasets.MNIST( os.path.join(args.out, 'data'), train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), # If input is normalized, input for WebDNN also have to be normalized with same parameter. # Default of datasets.MNIST is 0=black, 1=white # transforms.Normalize((0.1307,), (0.3081,)) ])), batch_size=64, shuffle=True) test_loader = torch.utils.data.DataLoader( datasets.MNIST( os.path.join(args.out, 'data'), train=False, transform=transforms.Compose([ transforms.ToTensor(), # transforms.Normalize((0.1307,), (0.3081,)) ])), batch_size=1000, shuffle=True) model = Net().to(device) optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5) for epoch in range(1, 4): train(args, model, device, train_loader, optimizer, epoch) test(args, model, device, test_loader) # export model as ONNX format dummy_input = torch.autograd.Variable(torch.randn(1, 1, 28, 28)) torch.onnx.export(model, dummy_input, model_path, verbose=True) # model conversion using WebDNN onnx_model = onnx.load(model_path) graph = ONNXConverter().convert(onnx_model) for backend in args.backend.split(","): exec_info = generate_descriptor(backend, graph) exec_info.save(args.out) # test data for demo output_test_samples( datasets.MNIST( os.path.join(args.out, 'data'), train=False, transform=transforms.Compose([ transforms.ToTensor(), # transforms.Normalize((0.1307,), (0.3081,)) ])), os.path.join(args.out, "test_samples.json"))
from webdnn.frontend.onnx import ONNXConverter import onnx from webdnn.backend import generate_descriptor graph = onnx.load('models/model.proto') print(graph) print('=' * 80) graph = ONNXConverter().convert(graph) print(graph) exec_info = generate_descriptor( "webassembly", graph) # also "webgpu", "webassembly", "webgl", "fallback" are available. exec_info.save("./output")
def template(N=2, H=5, W=5, Cin=7, Cout=9, KH=3, KW=3, SH=1, SW=1, PH=1, PW=1, DH=1, DW=1, use_bias=False, description: str = ""): x_shape = [N, Cin, H, W] w_shape = [Cout, Cin, KH, KW] vx = np.random.rand(*x_shape) vw = np.random.rand(*w_shape) vb = np.random.rand(Cout) if use_bias else None if DH != 1 or DW != 1: vy = chainer.functions.dilated_convolution_2d(vx, vw, b=vb, stride=[SH, SW], pad=[PH, PW], dilate=[DH, DW]).data else: vy = chainer.functions.convolution_2d(vx, vw, b=vb, stride=[SH, SW], pad=[PH, PW]).data x = make_tensor_value_info("x", x_shape) w = make_tensor_value_info("w", w_shape) y = make_tensor_value_info("y", vy.shape) kwargs = { "kernel_shape": [KH, KW], "strides": [SH, SW], "dilations": [DH, DW], "pads": [PH, PH, PW, PW] } if use_bias: b = make_tensor_value_info("b", vb.shape) operator = make_node("Conv", ["x", "w", "b"], ["y"], **kwargs) model = make_model([operator], [x, w, b], [y]) else: operator = make_node("Conv", ["x", "w"], ["y"], **kwargs) model = make_model([operator], [x, w], [y]) graph = ONNXConverter().convert(model) inputs = {graph.inputs[0]: vx, graph.inputs[1]: vw} if use_bias: inputs[graph.inputs[2]] = vb generate_kernel_test_case( description=f"[ONNX] Conv {description}", graph=graph, inputs=inputs, expected={graph.outputs[0]: vy}, )