mnist_pytorch.json

{
    "_model": {
        "inputs": ["image"],
        "outputs": ["metric"],

        "nodes": {
            "stem": {"operation": {"type": "_cell", "cell_name": "stem"}},
            "flatten": {"operation": {"type": "__torch__.torch.nn.Flatten"}},
            "fc1": {"operation": {"type": "__torch__.torch.nn.Linear", "parameters": {"out_features": 256, "in_features": 1024}}},
            "fc2": {"operation": {"type": "__torch__.torch.nn.Linear", "parameters": {"out_features": 10, "in_features": 256}}},
            "softmax": {"operation": {"type": "__torch__.torch.nn.Softmax"}}
        },

        "edges": [
            {"head": ["_inputs", 0], "tail": ["stem", null]},
            {"head": ["stem", null], "tail": ["flatten", null]},
            {"head": ["flatten", null], "tail": ["fc1", null]},
            {"head": ["fc1", null], "tail": ["fc2", null]},
            {"head": ["fc2", null], "tail": ["softmax", null]},
            {"head": ["softmax", null], "tail": ["_outputs", 0]}
        ]
    },

    "stem": {
        "nodes": {
            "conv1": {"operation": {"type": "__torch__.torch.nn.Conv2d", "parameters": {"out_channels": 32, "in_channels": 1, "kernel_size": 5}}},
            "pool1": {"operation": {"type": "__torch__.torch.nn.MaxPool2d", "parameters": {"kernel_size": 2}}},
            "conv2": {"operation": {"type": "__torch__.torch.nn.Conv2d", "parameters": {"out_channels": 64, "in_channels": 32, "kernel_size": 5}}},
            "pool2": {"operation": {"type": "__torch__.torch.nn.MaxPool2d", "parameters": {"kernel_size": 2}}}
        },

        "edges": [
            {"head": ["_inputs", 0], "tail": ["conv1", null]},
            {"head": ["conv1", null], "tail": ["pool1", null]},
            {"head": ["pool1", null], "tail": ["conv2", null]},
            {"head": ["conv2", null], "tail": ["pool2", null]},
            {"head": ["pool2", null], "tail": ["_outputs", 0]}
        ]
    },

    "_evaluator": {
        "type": "DebugEvaluator"
    }
}