def test_caffe2_simple_model(self): model = ModelHelper(name="mnist") # how come those inputs don't break the forward pass =.=a workspace.FeedBlob("data", np.random.randn(1, 3, 64, 64).astype(np.float32)) workspace.FeedBlob("label", np.random.randn(1, 1000).astype(np.int)) with core.NameScope("conv1"): conv1 = brew.conv(model, "data", 'conv1', dim_in=1, dim_out=20, kernel=5) # Image size: 24 x 24 -> 12 x 12 pool1 = brew.max_pool(model, conv1, 'pool1', kernel=2, stride=2) # Image size: 12 x 12 -> 8 x 8 conv2 = brew.conv(model, pool1, 'conv2', dim_in=20, dim_out=100, kernel=5) # Image size: 8 x 8 -> 4 x 4 pool2 = brew.max_pool(model, conv2, 'pool2', kernel=2, stride=2) with core.NameScope("classifier"): # 50 * 4 * 4 stands for dim_out from previous layer multiplied by the image size fc3 = brew.fc(model, pool2, 'fc3', dim_in=100 * 4 * 4, dim_out=500) relu = brew.relu(model, fc3, fc3) pred = brew.fc(model, relu, 'pred', 500, 10) softmax = brew.softmax(model, pred, 'softmax') xent = model.LabelCrossEntropy([softmax, "label"], 'xent') # compute the expected loss loss = model.AveragedLoss(xent, "loss") model.net.RunAllOnMKL() model.param_init_net.RunAllOnMKL() model.AddGradientOperators([loss], skip=1) blob_name_tracker = {} graph = c2_graph.model_to_graph_def( model, blob_name_tracker=blob_name_tracker, shapes={}, show_simplified=False, ) compare_proto(graph, self)
def test_simple_cnnmodel(self): model = cnn.CNNModelHelper("NCHW", name="overfeat") workspace.FeedBlob( "data", np.random.randn(1, 3, 64, 64).astype(np.float32)) workspace.FeedBlob("label", np.random.randn(1, 1000).astype(np.int)) with core.NameScope("conv1"): conv1 = model.Conv("data", "conv1", 3, 96, 11, stride=4) relu1 = model.Relu(conv1, conv1) pool1 = model.MaxPool(relu1, "pool1", kernel=2, stride=2) with core.NameScope("classifier"): fc = model.FC(pool1, "fc", 4096, 1000) pred = model.Softmax(fc, "pred") xent = model.LabelCrossEntropy([pred, "label"], "xent") loss = model.AveragedLoss(xent, "loss") blob_name_tracker = {} graph = c2_graph.model_to_graph_def( model, blob_name_tracker=blob_name_tracker, shapes={}, show_simplified=False, ) compare_proto(graph, self)
def add_graph(self, model, input_to_model=None, verbose=False): # prohibit second call? # no, let tensorboard handle it and show its warning message. torch._C._log_api_usage_once("tensorboard.logging.add_graph") if hasattr(model, 'forward'): # A valid PyTorch model should have a 'forward' method self._get_file_writer().add_graph( graph(model, input_to_model, verbose)) else: # Caffe2 models do not have the 'forward' method from caffe2.proto import caffe2_pb2 from caffe2.python import core from torch.utils.tensorboard._caffe2_graph import ( model_to_graph_def, nets_to_graph_def, protos_to_graph_def) if isinstance(model, list): if isinstance(model[0], core.Net): current_graph = nets_to_graph_def(model) elif isinstance(model[0], caffe2_pb2.NetDef): current_graph = protos_to_graph_def(model) else: # Handles cnn.CNNModelHelper, model_helper.ModelHelper current_graph = model_to_graph_def(model) event = event_pb2.Event( graph_def=current_graph.SerializeToString()) self._get_file_writer().add_event(event)