mindspore 1.10.0__cp37-cp37m-win_amd64.whl → 2.0.0rc1__cp37-cp37m-win_amd64.whl

mindspore/train/callback/_on_request_exit.py

@@ -0,0 +1,212 @@
+# Copyright 2022 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""OnRequestExit Callback class."""
+
+from __future__ import absolute_import
+import os
+import signal
+
+from mindspore import log
+from mindspore import _checkparam as Validator
+from mindspore.train.serialization import load_checkpoint, save_checkpoint, export
+from mindspore.train.callback._callback import Callback
+
+
+class OnRequestExit(Callback):
+    """
+    Respond to the user's closing request, exit the training or eval process, and save the checkpoint and mindir.
+
+    Register OnRequestExit Callback before training, when the user want to exit the training process
+    and save the training data, could send the registered exit signal 'sig' to the training process.
+    After the training process executes the current step, saves the current training status,
+    including checkpoint and mindir, and then exit the training process.
+
+    Args:
+        save_ckpt (bool): Whether save the checkpoint before the training process exit. Default: True.
+        save_mindir (bool): Whether save the mindir before the training process exit. Default: True.
+        file_name (str): The saved checkpoint and mindir file name,
+            the checkpoint file add suffix '.ckpt', the mindir file add suffix '.mindir'. Default: 'Net'.
+        directory (str): The directory save checkpoint and mindir. Default: './'.
+        sig (int): The user registered exit signal, it must be a captureable and negligible signal.
+            When the process receives the signal, exits the training or eval process. Default: signal.SIGTERM.
+
+    Raises:
+        ValueError: If the 'save_ckpt' is not a bool.
+        ValueError: If the 'save_mindir' is not a bool.
+        ValueError: If the 'file_name' is not a str.
+        ValueError: If the 'directory' is not a str.
+        ValueError: If the 'sig' is not an int or the 'sig' is signal.SIGKILL.
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import dataset as ds
+        >>> from mindspore import nn
+        >>>
+        >>> # Define the forward net
+        >>> class ForwardNet(nn.Cell):
+        >>>     def __init__(self, num_class=10, channel=1):
+        >>>         super(ForwardNet, self).__init__()
+        >>>         self.param = ms.Parameter(1.0)
+        >>>         self.relu = ms.ops.ReLU()
+        >>>
+        >>>     def construct(self, x):
+        >>>         return self.relu(x + self.param)
+        >>> forward_net = ForwardNet()
+        >>> loss = nn.MAELoss()
+        >>> opt = nn.Momentum(forward_net.trainable_params(), 0.01, 0.9)
+        >>> model = ms.Model(forward_net, loss_fn=loss, optimizer=opt)\
+        >>>
+        >>> # Create dataset
+        >>> def generator_multi_column():
+        >>>    i = 0
+        >>>    while i < 1000:
+        >>>        i += 1
+        >>>        yield np.ones((1, 32, 32)).astype(np.float32) * 0.01, np.array(1).astype(np.int32)
+        >>> dataset = ds.GeneratorDataset(source=generator_multi_column, column_names=["data", "label"])
+        >>> dataset = dataset.batch(32, drop_remainder=True)
+        >>>
+        >>> on_request_exit = ms.train.OnRequestExit(file_name='LeNet5')
+        >>> model.train(10, dataset, callbacks=on_request_exit)
+        >>> # The user send the signal SIGTERM to the training process,
+        >>> # the process would save the checkpoint and mindir, and then exit the training process.
+    """
+
+    def __init__(self, save_ckpt=True, save_mindir=True, file_name='Net', directory='./', sig=signal.SIGTERM):
+        super(OnRequestExit, self).__init__()
+        self.save_ckpt = Validator.check_isinstance('save_ckpt', save_ckpt, bool)
+        self.save_mindir = Validator.check_isinstance('save_mindir', save_mindir, bool)
+        if self.save_ckpt or self.save_mindir:
+            file_name = Validator.check_isinstance('file_name', file_name, str)
+            directory = Validator.check_isinstance('directory', directory, str)
+            os.makedirs(os.path.abspath(directory), exist_ok=True)
+            self.train_file_path = os.path.abspath(os.path.join(directory, f"{file_name}_train"))
+            self.eval_file_path = os.path.abspath(os.path.join(directory, f"{file_name}_eval"))
+        self.sig = Validator.check_isinstance('sig', sig, int)
+        if hasattr(signal, "SIGKILL") and self.sig == signal.SIGKILL:
+            raise ValueError("Not support send exit request by signal SIGKILL.")
+        self.exit = False
+
+    def on_train_begin(self, run_context):
+        """
+        When the train begin, register the handler for exit signal transferred by user.
+
+        Args:
+            run_context (RunContext): Context information of the model.
+                For more details, please refer to :class:`mindspore.train.RunContext`.
+        """
+        signal.signal(self.sig, self._handle_signal)
+        if self.save_ckpt and os.path.isfile(f"{self.train_file_path}.ckpt"):
+            cb_params = run_context.original_args()
+            train_net = cb_params.train_network
+            load_checkpoint(f"{self.train_file_path}.ckpt", net=train_net)
+
+    def on_train_step_end(self, run_context):
+        """
+        When the train step end, if received the exit signal, set the 'run_context' attribute '_stop_requested' to True.
+        Then exit the training process after this step training.
+
+        Args:
+            run_context (RunContext): Include some information of the model.
+                For more details, please refer to :class:`mindspore.train.RunContext`.
+        """
+        if self.exit:
+            run_context.request_stop()
+
+    def on_train_epoch_end(self, run_context):
+        """
+        When the train epoch end, if received the exit signal,
+        set the 'run_context' attribute '_stop_requested' to True.
+        Then exit the training process after this epoch training.
+
+        Args:
+            run_context (RunContext): Include some information of the model.
+                For more details, please refer to :class:`mindspore.train.RunContext`.
+        """
+        if self.exit:
+            run_context.request_stop()
+
+    def on_train_end(self, run_context):
+        """
+        When the train end, if received the exit signal,
+        the checkpoint and mindir would be saved according to the user config.
+
+        Args:
+            run_context (RunContext): Include some information of the model.
+                For more details, please refer to :class:`mindspore.train.RunContext`.
+        """
+        if not self.exit:
+            return
+        cb_params = run_context.original_args()
+        train_net = cb_params.train_network
+        if self.save_ckpt:
+            save_checkpoint(train_net, ckpt_file_name=self.train_file_path)
+        if self.save_mindir:
+            inputs = cb_params.train_dataset_element
+            export(train_net, *inputs, file_name=self.train_file_path, file_format='MINDIR')
+
+    def on_eval_begin(self, run_context):
+        """
+        When the eval begin, register the handler for exit signal transferred by user.
+
+        Args:
+            run_context (RunContext): Context information of the model.
+                For more details, please refer to :class:`mindspore.train.RunContext`.
+        """
+        signal.signal(self.sig, self._handle_signal)
+        if not self.save_ckpt:
+            return
+        cb_params = run_context.original_args()
+        eval_net = cb_params.eval_network
+        if os.path.isfile(f"{self.eval_file_path}.ckpt"):
+            load_checkpoint(f"{self.eval_file_path}.ckpt", net=eval_net)
+        elif os.path.isfile(f"{self.train_file_path}.ckpt"):
+            load_checkpoint(f"{self.train_file_path}.ckpt", net=eval_net)
+
+    def on_eval_step_end(self, run_context):
+        """
+        When the eval step end, if received the exit signal, set the 'run_context' attribute '_stop_requested' to True.
+        Then exit the eval process after this step eval.
+
+        Args:
+            run_context (RunContext): Include some information of the model.
+                For more details, please refer to :class:`mindspore.train.RunContext`.
+        """
+        if self.exit:
+            run_context.request_stop()
+
+    def on_eval_end(self, run_context):
+        """
+        When the eval end, if received the exit signal,
+        the checkpoint and mindir would be saved according to the user config.
+
+        Args:
+            run_context (RunContext): Include some information of the model.
+                For more details, please refer to :class:`mindspore.train.RunContext`.
+        """
+        if not self.exit:
+            return
+        cb_params = run_context.original_args()
+        eval_net = cb_params.eval_network
+        if self.save_ckpt:
+            save_checkpoint(eval_net, ckpt_file_name=self.eval_file_path)
+        if self.save_mindir:
+            inputs = cb_params.eval_dataset_element
+            export(eval_net, *inputs, file_name=self.eval_file_path, file_format='MINDIR')
+
+    def _handle_signal(self, signum, frame):
+        """Handle the received signal"""
+        log.debug(f"signum: {signum}, frame: {frame}")
+        self.exit = True

mindspore/train/callback/_reduce_lr_on_plateau.py

@@ -20,7 +20,7 @@ import numpy as np
 
 from mindspore.common.tensor import Tensor
 from mindspore.common.parameter import Parameter
-from mindspore._checkparam import Validator, Rel
+from mindspore import _checkparam as Validator
 from mindspore import log as logger
 from mindspore.ops import functional as F, ReduceOp
 from mindspore import nn, ops
@@ -81,8 +81,14 @@ class ReduceLROnPlateau(Callback):
         ValueError: The learning rate is not a Parameter.
 
     Examples:
-        >>> from mindspore.train.callback import ReduceLROnPlateau
-        >>> from mindspore import Model, nn
+        .. note::
+            Before running the following example, you need to customize the network LeNet5 and
+            dataset preparation function create_dataset. Refer to
+            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0/beginner/model.html>`_
+            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0/beginner/dataset.html>`_ .
+
+        >>> from mindspore import nn
+        >>> from mindspore.train import Model, ReduceLROnPlateau
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
@@ -96,7 +102,7 @@ class ReduceLROnPlateau(Callback):
                  mode='auto', min_delta=1e-4, cooldown=0, min_lr=0):
         super(ReduceLROnPlateau, self).__init__()
         self.monitor = Validator.check_value_type('monitor', monitor, str)
-        self.factor = Validator.check_float_range(factor, 0.0, 1.0, Rel.INC_NEITHER)
+        self.factor = Validator.check_float_range(factor, 0.0, 1.0, Validator.INC_NEITHER)
         self.patience = Validator.check_non_negative_int(patience)
         self.verbose = Validator.check_bool(verbose)
         self.mode = Validator.check_value_type('mode', mode, str)
@@ -124,7 +130,7 @@ class ReduceLROnPlateau(Callback):
 
         Args:
             run_context (RunContext): Context information of the model. For more details,
-                    please refer to :class:`mindspore.RunContext`.
+                    please refer to :class:`mindspore.train.RunContext`.
         """
         self.cooldown_counter = 0
         self.wait = 0
@@ -140,7 +146,7 @@ class ReduceLROnPlateau(Callback):
 
         Args:
             run_context (RunContext): Context information of the model. For more details,
-                    please refer to :class:`mindspore.RunContext`.
+                    please refer to :class:`mindspore.train.RunContext`.
         """
         cb_params = run_context.original_args()
         cur_lr = cb_params.optimizer.learning_rate
@@ -188,7 +194,7 @@ class ReduceLROnPlateau(Callback):
 
         Args:
             cb_params (dict): A dictionary stores context information of the model. For more details,
-                    please refer to :class:`mindspore.RunContext`.
+                    please refer to :class:`mindspore.train.RunContext`.
         """
         monitor_candidates = {}
         if self.monitor == "loss":

mindspore/train/callback/_summary_collector.py

@@ -103,7 +103,7 @@ class SummaryCollector(Callback):
               training computational graph is collected. Default: True.
             - collect_train_lineage (bool): Whether to collect lineage data for the training phase,
               this field will be displayed on the `lineage page \
-              <https://www.mindspore.cn/mindinsight/docs/en/r1.10/lineage_and_scalars_comparison.html>`_
+              <https://www.mindspore.cn/mindinsight/docs/en/r2.0/lineage_and_scalars_comparison.html>`_
               of MindInsight. Default: True.
             - collect_eval_lineage (bool): Whether to collect lineage data for the evaluation phase,
               this field will be displayed on the lineage page of MindInsight. Default: True.
@@ -177,6 +177,7 @@ class SummaryCollector(Callback):
     Examples:
         >>> import mindspore as ms
         >>> import mindspore.nn as nn
+        >>> from mindspore.train import Model, SummaryCollector
         >>> from mindspore.nn import Accuracy
         >>>
         >>> if __name__ == '__main__':
@@ -189,15 +190,15 @@ class SummaryCollector(Callback):
         ...     network = LeNet5(10)
         ...     net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
         ...     net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9)
-        ...     model = ms.Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}, amp_level="O2")
+        ...     model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}, amp_level="O2")
         ...
         ...     # Simple usage:
-        ...     summary_collector = ms.SummaryCollector(summary_dir='./summary_dir')
+        ...     summary_collector = SummaryCollector(summary_dir='./summary_dir')
         ...     model.train(1, ds_train, callbacks=[summary_collector], dataset_sink_mode=False)
         ...
         ...     # Do not collect metric and collect the first layer parameter, others are collected by default
         ...     specified={'collect_metric': False, 'histogram_regular': '^conv1.*'}
-        ...     summary_collector = ms.SummaryCollector(summary_dir='./summary_dir', collect_specified_data=specified)
+        ...     summary_collector = SummaryCollector(summary_dir='./summary_dir', collect_specified_data=specified)
         ...     model.train(1, ds_train, callbacks=[summary_collector], dataset_sink_mode=False)
     """
 
@@ -487,7 +488,26 @@ class SummaryCollector(Callback):
             self._record.add_value('custom_lineage_data', 'custom_lineage_data', packaged_custom_data)
             self._has_saved_custom_data = True
             self._record.record(cb_params.cur_step_num)
+        if not self._dataset_sink_mode:
+            self._collect_tensor_data(cb_params)
 
+        collect_landscape = self._collect_specified_data.get('collect_landscape')
+        if collect_landscape is not None:
+            intervals = collect_landscape.get('intervals')
+            collect_interval = False
+            for interval in intervals:
+                if "cur_step_num" in cb_params:
+                    if cb_params.cur_step_num in interval:
+                        collect_interval = True
+                        break
+
+        if collect_landscape and collect_landscape.get('unit', 'step') == 'step' and collect_interval:
+            self._save_model_params_for_landscape(cb_params)
+
+    def _collect_tensor_data(self, cb_params):
+        """Collect tensor summary data."""
+        if cb_params.mode != ModeEnum.TRAIN.value:
+            return
         if self._first_step:
             self._tensor_collect_range = self._get_tensor_collect_range(cb_params, self._dataset_sink_mode)
             self._collect_at_step_end(cb_params, plugin_filter=None)
@@ -502,19 +522,6 @@ class SummaryCollector(Callback):
             elif current % self._collect_freq == 0:
                 self._collect_at_step_end(cb_params, lambda plugin: plugin != PluginEnum.TENSOR.value)
 
-        collect_landscape = self._collect_specified_data.get('collect_landscape')
-        if collect_landscape is not None:
-            intervals = collect_landscape.get('intervals')
-            collect_interval = False
-            for interval in intervals:
-                if "cur_step_num" in cb_params:
-                    if cb_params.cur_step_num in interval:
-                        collect_interval = True
-                        break
-
-        if collect_landscape and collect_landscape.get('unit', 'step') == 'step' and collect_interval:
-            self._save_model_params_for_landscape(cb_params)
-
     def _get_tensor_collect_range(self, cb_params, dataset_sink_mode):
         """Get tensor collect range."""
         total_step = cb_params.epoch_num
@@ -541,6 +548,7 @@ class SummaryCollector(Callback):
 
     def epoch_end(self, run_context):
         cb_params = run_context.original_args()
+        self._collect_tensor_data(cb_params)
         collect_landscape = self._collect_specified_data.get('collect_landscape')
         if collect_landscape is not None:
             intervals = collect_landscape.get('intervals')
@@ -862,14 +870,14 @@ class SummaryCollector(Callback):
         if regular is not None:
             for parameter in parameters:
                 if re.match(regular, parameter.name):
-                    self._record.add_value(PluginEnum.HISTOGRAM.value, parameter.name+'/auto', parameter.data)
+                    self._record.add_value(PluginEnum.HISTOGRAM.value, parameter.name + '/auto', parameter.data)
             return
 
         # Note: If `histogram_regular` in `self._collect_specified_data` and the value is None,
         # we will collect the first five parameters.
         default_parameter_count = 5
         for parameter in parameters[:default_parameter_count]:
-            self._record.add_value(PluginEnum.HISTOGRAM.value, parameter.name+'/auto', parameter.data)
+            self._record.add_value(PluginEnum.HISTOGRAM.value, parameter.name + '/auto', parameter.data)
 
     @staticmethod
     def _get_learning_rate(optimizer):

mindspore/train/callback/_time_monitor.py

@@ -17,7 +17,7 @@ from __future__ import absolute_import
 
 import time
 
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.train.callback._callback import Callback
 
 
@@ -34,13 +34,19 @@ class TimeMonitor(Callback):
         ValueError: If data_size is not positive int.
 
     Examples:
-        >>> import mindspore as ms
+        .. note::
+            Before running the following example, you need to customize the network LeNet5 and
+            dataset preparation function create_dataset. Refer to
+            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0/beginner/model.html>`_
+            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0/beginner/dataset.html>`_ .
+
         >>> from mindspore import nn
+        >>> from mindspore.train import Model, TimeMonitor
         >>>
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
-        >>> model = ms.Model(net, loss_fn=loss, optimizer=optim)
+        >>> model = Model(net, loss_fn=loss, optimizer=optim)
         >>> data_path = './MNIST_Data'
         >>> dataset = create_dataset(data_path)
         >>> time_monitor = TimeMonitor()
@@ -57,8 +63,8 @@ class TimeMonitor(Callback):
         Record time at the beginning of epoch.
 
         Args:
-            run_context (RunContext): Context of the process running.  For more details,
-                    please refer to :class:`mindspore.RunContext`.
+            run_context (RunContext): Context of the process running. For more details,
+                    please refer to :class:`mindspore.train.RunContext`.
         """
         self.epoch_time = time.time()
 
@@ -67,8 +73,8 @@ class TimeMonitor(Callback):
         Print process cost time at the end of epoch.
 
         Args:
-           run_context (RunContext): Context of the process running.  For more details,
-                   please refer to :class:`mindspore.RunContext`.
+           run_context (RunContext): Context of the process running. For more details,
+                   please refer to :class:`mindspore.train.RunContext`.
         """
         epoch_seconds = (time.time() - self.epoch_time) * 1000
         step_size = self.data_size

mindspore/train/checkpoint_pb2.py

@@ -19,7 +19,7 @@ DESCRIPTOR = _descriptor.FileDescriptor(
   syntax='proto2',
   serialized_options=None,
   create_key=_descriptor._internal_create_key,
-  serialized_pb=b'\n\x10\x63heckpoint.proto\"b\n\nCheckpoint\x12 \n\x05value\x18\x01 \x03(\x0b\x32\x11.Checkpoint.Value\x1a\x32\n\x05Value\x12\x0b\n\x03tag\x18\x01 \x02(\t\x12\x1c\n\x06tensor\x18\x02 \x02(\x0b\x32\x0c.TensorProto\"H\n\x0bTensorProto\x12\x0c\n\x04\x64ims\x18\x01 \x03(\x03\x12\x13\n\x0btensor_type\x18\x02 \x02(\t\x12\x16\n\x0etensor_content\x18\x03 \x02(\x0c'
+  serialized_pb=b'\n\x10\x63heckpoint.proto\"\x93\x01\n\nCheckpoint\x12 \n\x05value\x18\x01 \x03(\x0b\x32\x11.Checkpoint.Value\x1a\x63\n\x05Value\x12\x0b\n\x03tag\x18\x01 \x02(\t\x12\x1e\n\x06tensor\x18\x02 \x01(\x0b\x32\x0c.TensorProtoH\x00\x12$\n\tmaptensor\x18\x03 \x01(\x0b\x32\x0f.MapTensorProtoH\x00\x42\x07\n\x05value\"H\n\x0bTensorProto\x12\x0c\n\x04\x64ims\x18\x01 \x03(\x03\x12\x13\n\x0btensor_type\x18\x02 \x02(\t\x12\x16\n\x0etensor_content\x18\x03 \x02(\x0c\".\n\x0eMapTensorProto\x12\x1c\n\x06tensor\x18\x01 \x03(\x0b\x32\x0c.TensorProto'
 )
 
 
@@ -42,7 +42,14 @@ _CHECKPOINT_VALUE = _descriptor.Descriptor(
       serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
     _descriptor.FieldDescriptor(
       name='tensor', full_name='Checkpoint.Value.tensor', index=1,
-      number=2, type=11, cpp_type=10, label=2,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='maptensor', full_name='Checkpoint.Value.maptensor', index=2,
+      number=3, type=11, cpp_type=10, label=1,
       has_default_value=False, default_value=None,
       message_type=None, enum_type=None, containing_type=None,
       is_extension=False, extension_scope=None,
@@ -58,9 +65,14 @@ _CHECKPOINT_VALUE = _descriptor.Descriptor(
   syntax='proto2',
   extension_ranges=[],
   oneofs=[
+    _descriptor.OneofDescriptor(
+      name='value', full_name='Checkpoint.Value.value',
+      index=0, containing_type=None,
+      create_key=_descriptor._internal_create_key,
+    fields=[]),
   ],
-  serialized_start=68,
-  serialized_end=118,
+  serialized_start=69,
+  serialized_end=168,
 )
 
 _CHECKPOINT = _descriptor.Descriptor(
@@ -90,8 +102,8 @@ _CHECKPOINT = _descriptor.Descriptor(
   extension_ranges=[],
   oneofs=[
   ],
-  serialized_start=20,
-  serialized_end=118,
+  serialized_start=21,
+  serialized_end=168,
 )
 
 
@@ -136,15 +148,56 @@ _TENSORPROTO = _descriptor.Descriptor(
   extension_ranges=[],
   oneofs=[
   ],
-  serialized_start=120,
-  serialized_end=192,
+  serialized_start=170,
+  serialized_end=242,
+)
+
+
+_MAPTENSORPROTO = _descriptor.Descriptor(
+  name='MapTensorProto',
+  full_name='MapTensorProto',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='tensor', full_name='MapTensorProto.tensor', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=244,
+  serialized_end=290,
 )
 
 _CHECKPOINT_VALUE.fields_by_name['tensor'].message_type = _TENSORPROTO
+_CHECKPOINT_VALUE.fields_by_name['maptensor'].message_type = _MAPTENSORPROTO
 _CHECKPOINT_VALUE.containing_type = _CHECKPOINT
+_CHECKPOINT_VALUE.oneofs_by_name['value'].fields.append(
+  _CHECKPOINT_VALUE.fields_by_name['tensor'])
+_CHECKPOINT_VALUE.fields_by_name['tensor'].containing_oneof = _CHECKPOINT_VALUE.oneofs_by_name['value']
+_CHECKPOINT_VALUE.oneofs_by_name['value'].fields.append(
+  _CHECKPOINT_VALUE.fields_by_name['maptensor'])
+_CHECKPOINT_VALUE.fields_by_name['maptensor'].containing_oneof = _CHECKPOINT_VALUE.oneofs_by_name['value']
 _CHECKPOINT.fields_by_name['value'].message_type = _CHECKPOINT_VALUE
+_MAPTENSORPROTO.fields_by_name['tensor'].message_type = _TENSORPROTO
 DESCRIPTOR.message_types_by_name['Checkpoint'] = _CHECKPOINT
 DESCRIPTOR.message_types_by_name['TensorProto'] = _TENSORPROTO
+DESCRIPTOR.message_types_by_name['MapTensorProto'] = _MAPTENSORPROTO
 _sym_db.RegisterFileDescriptor(DESCRIPTOR)
 
 Checkpoint = _reflection.GeneratedProtocolMessageType('Checkpoint', (_message.Message,), {
@@ -169,5 +222,12 @@ TensorProto = _reflection.GeneratedProtocolMessageType('TensorProto', (_message.
   })
 _sym_db.RegisterMessage(TensorProto)
 
+MapTensorProto = _reflection.GeneratedProtocolMessageType('MapTensorProto', (_message.Message,), {
+  'DESCRIPTOR' : _MAPTENSORPROTO,
+  '__module__' : 'checkpoint_pb2'
+  # @@protoc_insertion_point(class_scope:MapTensorProto)
+  })
+_sym_db.RegisterMessage(MapTensorProto)
+
 
 # @@protoc_insertion_point(module_scope)