• Search by Module
• Search by Words
• Search Projects

• Java
• Python
• JavaScript
• TypeScript
• C++
• Scala
• Blog

• deep500-master
  • samples
    • accelegrad.py
    • ucf101_resnetlstm.py
    • train_until.py
    • tf_native_training.py
    • mnist_bias.py
    • network_inference.py
    • onnxtest.py
    • synthetic_dataset.py
    • train_mnist.py
    • dataset_latency.py
    • distributed_training.py
    • deepbench
      • gemm_overhead_tensorflow.py
      • conv_overhead_caffe2.py
      • deepbench.py
      • gemm_overhead_pytorch.py
      • __init__.py
      • conv_overhead_pytorch.py
      • conv_overhead_tensorflow.py
      • gemm_overhead_caffe2.py
    • tf_native_operator.py
  • LICENSE
  • recipes
    • cifar10_resnet44.py
    • __init__.py
    • imagenet_resnet50.py
  • deep500.svg
  • setup.py
  • tutorials
    • operator.md
    • custom_optimizer.ipynb
    • README.md
  • deep500
    • recipe.py
    • datasets
      • mnist.py
      • ucf101.py
      • __init__.py
      • imagenet.py
      • cifar.py
    • utils
      • metric.py
      • mpi_helper.py
      • download.py
      • metrics.py
      • tensor_desc.py
      • event.py
      • __init__.py
      • onnx_interop
        • generated_operators.py
        • extended_operations_visitor.py
        • generate_operators.py
        • onnx_base_visitor.py
        • extended_operations.py
        • losses.py
        • loss_operations_visitor.py
        • onnx_objects.py
        • __init__.py
        • parser.py
        • operations_visitor.py
        • emptyoperations_visitor.py
        • onnx_test_parser.py
      • device.py
    • lv0
      • operators
        • operator_interface.py
        • include
          • deep500
            • deep500.h
        • cmake_wrapper.py
        • op_compiler.py
        • __init__.py
        • tools
          • get_cuda_arch.cpp
      • __init__.py
      • validation
        • metrics.py
        • forward.py
        • gradient.py
        • __init__.py
    • lv1
      • graph_executor.py
      • event.py
      • __init__.py
      • validation
        • executor.py
        • __init__.py
      • network.py
    • networks
      • pytorch_wide_resnet.py
      • pytorch_mnist.py
      • __init__.py
      • pytorch_resnet.py
      • pytorch_deep_mnist.py
    • lv2
      • summaries.py
      • optimizer.py
      • sampler.py
      • event.py
      • __init__.py
      • validation
        • optimizer.py
        • sampler.py
        • training.py
        • metrics.py
        • __init__.py
      • dataset.py
      • trainer.py
      • events
        • save_wrongs.py
        • summary_lossplot.py
        • stopping_conditions.py
        • training_progress.py
        • summary_boxplots.py
        • hyperparameter_schedule.py
        • summary_generator.py
        • summary_csv.py
        • __init__.py
    • frameworks
      • caffe2
        • support
          • mpi_operations.cc
          • mpi_operations.h
          • CMakeLists.txt
        • caffe2_graph_executor.py
        • caffe2_network.py
        • __init__.py
        • op_validation.py
        • caffe2_visitor_impl.py
        • caffe2_distributed_optimizer.py
      • pytorch
        • pytorch_network.py
        • pytorch_graph_executor.py
        • custom_operators
          • pytorch.tmpl.cpp
          • pytorch.py
          • CMakeLists.txt
          • __init__.py
        • pytorch_optimizers.py
        • __init__.py
        • op_validation.py
        • pytorch_visitor.py
      • reference
        • optimizers.py
        • reference_build_graph_visitor_impl_cpp.py
        • custom_operators
          • cpp
            • utility
              • tensor
                • transpose2d.hpp
                • broadcast.hpp
                • replicate.hpp
                • set_all.hpp
                • indexing.hpp
                • copy_data.hpp
              • math
                • convolution
                  • crosscorrelation_swap_axes.hpp
                  • crosscorrelation.hpp
                  • convolution_output_shape.hpp
                  • crosscorrelation_dilx_flipw.hpp
                  • autopadding.hpp
                • gemm.hpp
                • variance.hpp
                • mean.hpp
            • operators
              • add_op.cpp
              • softmax_op.cpp
              • gemm_op.cpp
              • matmul_op.cpp
              • averagepool_op.cpp
              • sum_op.cpp
              • conv_op.cpp
              • batchnormalization_op_inference.cpp
              • globalaveragepool_op.cpp
              • reshape_op.cpp
              • maxpool_op.cpp
              • relu_op.cpp
              • batchnormalization_op_training.cpp
          • custom_op_python.py
          • python
            • crossentropy_op.py
            • softmax_crossentropy_op.py
            • matmul_op.py
            • squeeze_op.py
            • sum_op.py
            • add_op.py
            • gemm_op.py
            • meansquarederror_op.py
            • pool_op_common.py
            • conv_op_common.py
            • batchnormalization_op_inference.py
            • conv_op.py
            • maxpool_op.py
            • __init__.py
            • relu_op.py
            • averagepool_op.py
            • globalaveragepool_op.py
            • reshape_op.py
            • batchnormalization_op_training.py
            • unsqueeze_op.py
            • softmax_op.py
          • base.tmpl.cpp
          • CMakeLists.txt
          • __init__.py
          • custom_op_cpp.py
          • base.py
        • augmentation.py
        • reference_graph_executor.py
        • reference_build_graph_visitor_impl.py
        • __init__.py
        • reference_network.py
        • distributed_optimizers.py
      • __init__.py
      • tensorflow
        • tf_optimizers.py
        • custom_operators
          • tf.tmpl.cpp
          • CMakeLists.txt
          • __init__.py
          • tf.py
        • tf_graph_executor.py
        • tf_distributed_optimizer.py
        • __init__.py
        • tf_network.py
        • op_validation.py
        • tf_visitor_impl.py
    • lv3
      • metrics.py
      • distributed_optimizer.py
      • distributed_sampler.py
      • __init__.py
      • communication.py
    • __init__.py
  • README.md
  • .gitignore

import os
from typing import List

from caffe2.proto import caffe2_pb2
from caffe2.proto.caffe2_pb2 import DeviceOption
from caffe2.python.modeling import initializers
from caffe2.python.predictor import mobile_exporter

from caffe2.python import core, workspace, model_helper, net_drawer

from deep500.lv1.network import Network


class Caffe2Network(Network):
    def __init__(self, device_option: DeviceOption):
        super(Caffe2Network, self).__init__()
        self.device_option = device_option

        self.train_model = model_helper.ModelHelper(name="train_default_net")
        self.test_model = model_helper.ModelHelper(name="test_default_net", init_params=False)
        self.train_net = self.train_model.net
        self.test_net = self.test_model.net
        self.train_init_net = self.train_model.param_init_net
        self.test_init_net = self.test_model.param_init_net
        self.workspace = workspace
        self.output_dict = {}
        self.param_names = None
        # dict that helps us remember that we already added the gradients to the graph for a given loss
        self.gradients_by_loss = {}
        self.is_cuda = (device_option.device_type == caffe2_pb2.CUDA)

    def _setup(self):
        self.workspace.RunNetOnce(self.train_init_net)
        self.workspace.RunNetOnce(self.test_init_net)

    def get_params(self):
        if self.param_names is None:
            self.param_names = [p._name for p in self.train_model.params]
        return self.param_names

    def _teardown(self):
        workspace.ResetWorkspace()

    def add_loss(self, y: str):
        if y != '' and y not in self.gradients_by_loss:
            self.gradients_by_loss[y] = y
            self.train_model.AddGradientOperators([y])

    def get_input_nodes(self):
        return list(self.input_dict.keys())
            
    def get_output_nodes(self):
        return list(self.output_dict.keys())

    def add_optimizer(self, optimizer):
        raise Exception('we do not use this')

    def gradient(self, y: str = 'loss'):
        gradients = self.train_model.param_to_grad.items()
        return gradients

    def _custom_parsing_context(self):
        return core.DeviceScope(self.device_option)

    def export_native(self, folder: str, inputs: List[str]):
        print("Save the model to init_net.pb and predict_net.pb")

        _, init_net2 = mobile_exporter.Export(workspace, self.train_model.param_init_net, [])
        init_net, predict_net = mobile_exporter.Export(workspace, self.train_model.net, self.train_model.params)
        if not os.path.exists(folder):
            os.makedirs(folder)

        with open("{}/train_init_net2.pb".format(folder), 'wb') as f:
            f.write(init_net2.SerializeToString())

        with open("{}/train_init_net.pb".format(folder), 'wb') as f:
            f.write(init_net.SerializeToString())

        with open("{}/train_predict_net.pb".format(folder), 'wb') as f:
            f.write(predict_net.SerializeToString())
        print("Saved the model to folder: {} and names train_predict_net.pb and train_init_net.pb".format(folder))

        init_net, predict_net = mobile_exporter.Export(workspace, self.test_model.net, self.test_model.params)
        with open("{}/test_init_net.pb".format(folder), 'wb') as f:
            f.write(init_net.SerializeToString())

        with open("{}/test_predict_net.pb".format(folder), 'wb') as f:
            f.write(predict_net.SerializeToString())
        print("Saved the model to folder: {} and names test_predict_net.pb and test_init_net.pb".format(folder))

    def teardown(self):
        pass

    def add_output(self, output: str):
        self.output_dict[output] = None

    def fetch_tensors(self, names):
        return [workspace.FetchBlob(name) for name in names]

    def feed_tensor(self, name, new_value, device_option=None, is_param=False):
        if is_param:
            self.train_model.create_param(name, new_value.shape,
                                          initializer=initializers.ExternalInitializer())
        device_option = device_option if device_option else self.device_option
        return workspace.FeedBlob(name, new_value, device_option)