mindspore 2.1.0__cp39-cp39-win_amd64.whl → 2.2.10__cp39-cp39-win_amd64.whl

mindspore/train/anf_ir_pb2.py

@@ -20,7 +20,7 @@ DESCRIPTOR = _descriptor.FileDescriptor(
   syntax='proto2',
   serialized_options=None,
   create_key=_descriptor._internal_create_key,
-  serialized_pb=b'\n\x0c\x61nf_ir.proto\x12\x0emindspore.irpb\"\xdb\x04\n\nValueProto\x12\'\n\x05\x64type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x10\n\x08\x62ool_val\x18\x02 \x01(\x08\x12\x0f\n\x07int_val\x18\x03 \x01(\x03\x12\x10\n\x08uint_val\x18\x04 \x01(\x04\x12\x11\n\tfloat_val\x18\x05 \x01(\x02\x12\x12\n\ndouble_val\x18\x06 \x01(\x01\x12\x0f\n\x07str_val\x18\x07 \x01(\t\x12/\n\ntensor_val\x18\x08 \x01(\x0b\x32\x1b.mindspore.irpb.TensorProto\x12)\n\x05graph\x18\t \x01(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12\x11\n\tbool_vals\x18\n \x03(\x08\x12\x10\n\x08int_vals\x18\x0b \x03(\x03\x12\x11\n\tuint_vals\x18\x0c \x03(\x04\x12\x12\n\nfloat_vals\x18\r \x03(\x02\x12\x13\n\x0b\x64ouble_vals\x18\x0e \x03(\x01\x12\x10\n\x08str_vals\x18\x0f \x03(\t\x12\x30\n\x0btensor_vals\x18\x10 \x03(\x0b\x32\x1b.mindspore.irpb.TensorProto\x12*\n\x06graphs\x18\x11 \x03(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12*\n\x06values\x18\x12 \x03(\x0b\x32\x1a.mindspore.irpb.ValueProto\x12\x31\n\x08\x64ict_val\x18\x13 \x03(\x0b\x32\x1f.mindspore.irpb.NamedValueProto\x12+\n\x08type_val\x18\x14 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\"I\n\x0e\x41ttributeProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"I\n\x0fNamedValueProto\x12\x0b\n\x03key\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"t\n\x10TensorShapeProto\x12\x37\n\x03\x64im\x18\x01 \x03(\x0b\x32*.mindspore.irpb.TensorShapeProto.Dimension\x1a\'\n\tDimension\x12\x0c\n\x04size\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\t\"\xda\x02\n\tTypeProto\x12+\n\tdata_type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x37\n\x0btensor_type\x18\x02 \x01(\x0b\x32 .mindspore.irpb.TypeProto.TensorH\x00\x12;\n\rsequence_type\x18\x03 \x01(\x0b\x32\".mindspore.irpb.TypeProto.SequenceH\x00\x1a\x66\n\x06Tensor\x12+\n\telem_type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12/\n\x05shape\x18\x02 \x01(\x0b\x32 .mindspore.irpb.TensorShapeProto\x1a\x39\n\x08Sequence\x12-\n\nelem_types\x18\x01 \x03(\x0b\x32\x19.mindspore.irpb.TypeProtoB\x07\n\x05value\"x\n\x0eParameterProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\'\n\x04type\x18\x02 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\x12/\n\x0b\x64\x65\x66\x61ult_val\x18\x03 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"D\n\x0bOutputProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\'\n\x04type\x18\x02 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\"z\n\nInputProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x31\n\x04type\x18\x02 \x01(\x0e\x32#.mindspore.irpb.InputProto.EdgeType\"+\n\x08\x45\x64geType\x12\r\n\tDATA_EDGE\x10\x00\x12\x10\n\x0c\x43ONTROL_EDGE\x10\x01\"\x83\x02\n\tNodeProto\x12)\n\x05input\x18\x01 \x03(\x0b\x32\x1a.mindspore.irpb.InputProto\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0f\n\x07op_type\x18\x03 \x01(\t\x12\r\n\x05scope\x18\x04 \x01(\t\x12\x31\n\tattribute\x18\x05 \x03(\x0b\x32\x1e.mindspore.irpb.AttributeProto\x12.\n\x0boutput_type\x18\x06 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\x12\x10\n\x08output_i\x18\x07 \x01(\x04\x12\x11\n\tfull_name\x18\x08 \x01(\t\x12\x15\n\rinstance_name\x18\n \x01(\t\"\xb0\x01\n\nModelProto\x12\x12\n\nir_version\x18\x01 \x01(\x03\x12\x0e\n\x06\x64omain\x18\x02 \x01(\t\x12\x15\n\rmodel_version\x18\x03 \x01(\x03\x12)\n\x05graph\x18\x04 \x01(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12<\n\x12metadata_operators\x18\x05 \x01(\x0b\x32 .mindspore.irpb.OperatorSetProto\"?\n\rOperatorProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06\x63onfig\x18\x02 \x01(\x0c\x12\x10\n\x08obj_info\x18\x03 \x01(\x0c\"U\n\x10OperatorSetProto\x12\x30\n\toperators\x18\x01 \x03(\x0b\x32\x1d.mindspore.irpb.OperatorProto\x12\x0f\n\x07summary\x18\x02 \x01(\t\"\xda\x01\n\nGraphProto\x12\'\n\x04node\x18\x01 \x03(\x0b\x32\x19.mindspore.irpb.NodeProto\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x32\n\nparameters\x18\x03 \x03(\x0b\x32\x1e.mindspore.irpb.ParameterProto\x12,\n\x07outputs\x18\x04 \x03(\x0b\x32\x1b.mindspore.irpb.OutputProto\x12\x33\n\nconst_vals\x18\x05 \x03(\x0b\x32\x1f.mindspore.irpb.NamedValueProto\"\xd4\x01\n\x0bTensorProto\x12\x0c\n\x04\x64ims\x18\x01 \x03(\x03\x12+\n\tdata_type\x18\x02 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x16\n\nfloat_data\x18\x03 \x03(\x02\x42\x02\x10\x01\x12\x16\n\nint32_data\x18\x04 \x03(\x05\x42\x02\x10\x01\x12\x16\n\nint64_data\x18\x05 \x03(\x03\x42\x02\x10\x01\x12\x17\n\x0b\x64ouble_data\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x17\n\x0buint64_data\x18\x07 \x03(\x04\x42\x02\x10\x01\x12\x10\n\x08raw_data\x18\x08 \x01(\x0c*/\n\x07Version\x12\x14\n\x10UNKNOWWN_VERSION\x10\x00\x12\x0e\n\nIR_VERSION\x10\x01*\xcb\x05\n\x08\x44\x61taType\x12\x10\n\x0c\x44T_UNDEFINED\x10\x00\x12\x0b\n\x07\x44T_BOOL\x10\x01\x12\x0b\n\x07\x44T_INT8\x10\x02\x12\x0c\n\x08\x44T_INT16\x10\x03\x12\x0c\n\x08\x44T_INT32\x10\x04\x12\x0c\n\x08\x44T_INT64\x10\x05\x12\x0c\n\x08\x44T_UINT8\x10\x06\x12\r\n\tDT_UINT16\x10\x07\x12\r\n\tDT_UINT32\x10\x08\x12\r\n\tDT_UINT64\x10\t\x12\x0e\n\nDT_FLOAT16\x10\n\x12\x0e\n\nDT_FLOAT32\x10\x0b\x12\x0e\n\nDT_FLOAT64\x10\x0c\x12\r\n\tDT_STRING\x10\r\x12\r\n\tDT_TENSOR\x10\x0e\x12\x0c\n\x08\x44T_GRAPH\x10\x0f\x12\x0c\n\x08\x44T_BOOLS\x10\x10\x12\x0c\n\x08\x44T_INTS8\x10\x11\x12\r\n\tDT_INTS16\x10\x12\x12\r\n\tDT_INTS32\x10\x13\x12\r\n\tDT_INTS64\x10\x14\x12\r\n\tDT_UINTS8\x10\x15\x12\x0e\n\nDT_UINTS16\x10\x16\x12\x0e\n\nDT_UINTS32\x10\x17\x12\x0e\n\nDT_UINTS64\x10\x18\x12\x0f\n\x0b\x44T_FLOATS16\x10\x19\x12\x0f\n\x0b\x44T_FLOATS32\x10\x1a\x12\x0f\n\x0b\x44T_FLOATS64\x10\x1b\x12\x0e\n\nDT_STRINGS\x10\x1c\x12\x0e\n\nDT_TENSORS\x10\x1d\x12\r\n\tDT_GRAPHS\x10\x1e\x12\x0c\n\x08\x44T_TUPLE\x10\x1f\x12\x0b\n\x07\x44T_LIST\x10 \x12\x0b\n\x07\x44T_DICT\x10!\x12\x0b\n\x07\x44T_NONE\x10\"\x12\x0f\n\x0b\x44T_SYM_INST\x10#\x12\x0f\n\x0b\x44T_BASE_INT\x10$\x12\x10\n\x0c\x44T_BASE_UINT\x10%\x12\x11\n\rDT_BASE_FLOAT\x10&\x12\x0b\n\x07\x44T_TYPE\x10\'\x12\n\n\x06\x44T_ANY\x10(\x12\r\n\tDT_REFKEY\x10)\x12\n\n\x06\x44T_REF\x10*\x12\x10\n\x0c\x44T_COMPLEX64\x10+\x12\x11\n\rDT_COMPLEX128\x10,\x12\x13\n\x0f\x44T_BASE_COMPLEX\x10-'
+  serialized_pb=b'\n\x0c\x61nf_ir.proto\x12\x0emindspore.irpb\"\xdb\x04\n\nValueProto\x12\'\n\x05\x64type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x10\n\x08\x62ool_val\x18\x02 \x01(\x08\x12\x0f\n\x07int_val\x18\x03 \x01(\x03\x12\x10\n\x08uint_val\x18\x04 \x01(\x04\x12\x11\n\tfloat_val\x18\x05 \x01(\x02\x12\x12\n\ndouble_val\x18\x06 \x01(\x01\x12\x0f\n\x07str_val\x18\x07 \x01(\t\x12/\n\ntensor_val\x18\x08 \x01(\x0b\x32\x1b.mindspore.irpb.TensorProto\x12)\n\x05graph\x18\t \x01(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12\x11\n\tbool_vals\x18\n \x03(\x08\x12\x10\n\x08int_vals\x18\x0b \x03(\x03\x12\x11\n\tuint_vals\x18\x0c \x03(\x04\x12\x12\n\nfloat_vals\x18\r \x03(\x02\x12\x13\n\x0b\x64ouble_vals\x18\x0e \x03(\x01\x12\x10\n\x08str_vals\x18\x0f \x03(\t\x12\x30\n\x0btensor_vals\x18\x10 \x03(\x0b\x32\x1b.mindspore.irpb.TensorProto\x12*\n\x06graphs\x18\x11 \x03(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12*\n\x06values\x18\x12 \x03(\x0b\x32\x1a.mindspore.irpb.ValueProto\x12\x31\n\x08\x64ict_val\x18\x13 \x03(\x0b\x32\x1f.mindspore.irpb.NamedValueProto\x12+\n\x08type_val\x18\x14 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\"I\n\x0e\x41ttributeProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"I\n\x0fNamedValueProto\x12\x0b\n\x03key\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"t\n\x10TensorShapeProto\x12\x37\n\x03\x64im\x18\x01 \x03(\x0b\x32*.mindspore.irpb.TensorShapeProto.Dimension\x1a\'\n\tDimension\x12\x0c\n\x04size\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\t\"\xda\x02\n\tTypeProto\x12+\n\tdata_type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x37\n\x0btensor_type\x18\x02 \x01(\x0b\x32 .mindspore.irpb.TypeProto.TensorH\x00\x12;\n\rsequence_type\x18\x03 \x01(\x0b\x32\".mindspore.irpb.TypeProto.SequenceH\x00\x1a\x66\n\x06Tensor\x12+\n\telem_type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12/\n\x05shape\x18\x02 \x01(\x0b\x32 .mindspore.irpb.TensorShapeProto\x1a\x39\n\x08Sequence\x12-\n\nelem_types\x18\x01 \x03(\x0b\x32\x19.mindspore.irpb.TypeProtoB\x07\n\x05value\"x\n\x0eParameterProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\'\n\x04type\x18\x02 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\x12/\n\x0b\x64\x65\x66\x61ult_val\x18\x03 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"D\n\x0bOutputProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\'\n\x04type\x18\x02 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\"z\n\nInputProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x31\n\x04type\x18\x02 \x01(\x0e\x32#.mindspore.irpb.InputProto.EdgeType\"+\n\x08\x45\x64geType\x12\r\n\tDATA_EDGE\x10\x00\x12\x10\n\x0c\x43ONTROL_EDGE\x10\x01\"\x83\x02\n\tNodeProto\x12)\n\x05input\x18\x01 \x03(\x0b\x32\x1a.mindspore.irpb.InputProto\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0f\n\x07op_type\x18\x03 \x01(\t\x12\r\n\x05scope\x18\x04 \x01(\t\x12\x31\n\tattribute\x18\x05 \x03(\x0b\x32\x1e.mindspore.irpb.AttributeProto\x12.\n\x0boutput_type\x18\x06 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\x12\x10\n\x08output_i\x18\x07 \x01(\x04\x12\x11\n\tfull_name\x18\x08 \x01(\t\x12\x15\n\rinstance_name\x18\n \x01(\t\"\xb0\x01\n\nModelProto\x12\x12\n\nir_version\x18\x01 \x01(\x03\x12\x0e\n\x06\x64omain\x18\x02 \x01(\t\x12\x15\n\rmodel_version\x18\x03 \x01(\x03\x12)\n\x05graph\x18\x04 \x01(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12<\n\x12metadata_operators\x18\x05 \x01(\x0b\x32 .mindspore.irpb.OperatorSetProto\"?\n\rOperatorProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06\x63onfig\x18\x02 \x01(\x0c\x12\x10\n\x08obj_info\x18\x03 \x01(\x0c\"U\n\x10OperatorSetProto\x12\x30\n\toperators\x18\x01 \x03(\x0b\x32\x1d.mindspore.irpb.OperatorProto\x12\x0f\n\x07summary\x18\x02 \x01(\t\"\xda\x01\n\nGraphProto\x12\'\n\x04node\x18\x01 \x03(\x0b\x32\x19.mindspore.irpb.NodeProto\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x32\n\nparameters\x18\x03 \x03(\x0b\x32\x1e.mindspore.irpb.ParameterProto\x12,\n\x07outputs\x18\x04 \x03(\x0b\x32\x1b.mindspore.irpb.OutputProto\x12\x33\n\nconst_vals\x18\x05 \x03(\x0b\x32\x1f.mindspore.irpb.NamedValueProto\"\xd4\x01\n\x0bTensorProto\x12\x0c\n\x04\x64ims\x18\x01 \x03(\x03\x12+\n\tdata_type\x18\x02 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x16\n\nfloat_data\x18\x03 \x03(\x02\x42\x02\x10\x01\x12\x16\n\nint32_data\x18\x04 \x03(\x05\x42\x02\x10\x01\x12\x16\n\nint64_data\x18\x05 \x03(\x03\x42\x02\x10\x01\x12\x17\n\x0b\x64ouble_data\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x17\n\x0buint64_data\x18\x07 \x03(\x04\x42\x02\x10\x01\x12\x10\n\x08raw_data\x18\x08 \x01(\x0c*/\n\x07Version\x12\x14\n\x10UNKNOWWN_VERSION\x10\x00\x12\x0e\n\nIR_VERSION\x10\x01*\xee\x05\n\x08\x44\x61taType\x12\x10\n\x0c\x44T_UNDEFINED\x10\x00\x12\x0b\n\x07\x44T_BOOL\x10\x01\x12\x0b\n\x07\x44T_INT8\x10\x02\x12\x0c\n\x08\x44T_INT16\x10\x03\x12\x0c\n\x08\x44T_INT32\x10\x04\x12\x0c\n\x08\x44T_INT64\x10\x05\x12\x0c\n\x08\x44T_UINT8\x10\x06\x12\r\n\tDT_UINT16\x10\x07\x12\r\n\tDT_UINT32\x10\x08\x12\r\n\tDT_UINT64\x10\t\x12\x0e\n\nDT_FLOAT16\x10\n\x12\x0e\n\nDT_FLOAT32\x10\x0b\x12\x0e\n\nDT_FLOAT64\x10\x0c\x12\r\n\tDT_STRING\x10\r\x12\r\n\tDT_TENSOR\x10\x0e\x12\x0c\n\x08\x44T_GRAPH\x10\x0f\x12\x0c\n\x08\x44T_BOOLS\x10\x10\x12\x0c\n\x08\x44T_INTS8\x10\x11\x12\r\n\tDT_INTS16\x10\x12\x12\r\n\tDT_INTS32\x10\x13\x12\r\n\tDT_INTS64\x10\x14\x12\r\n\tDT_UINTS8\x10\x15\x12\x0e\n\nDT_UINTS16\x10\x16\x12\x0e\n\nDT_UINTS32\x10\x17\x12\x0e\n\nDT_UINTS64\x10\x18\x12\x0f\n\x0b\x44T_FLOATS16\x10\x19\x12\x0f\n\x0b\x44T_FLOATS32\x10\x1a\x12\x0f\n\x0b\x44T_FLOATS64\x10\x1b\x12\x0e\n\nDT_STRINGS\x10\x1c\x12\x0e\n\nDT_TENSORS\x10\x1d\x12\r\n\tDT_GRAPHS\x10\x1e\x12\x0c\n\x08\x44T_TUPLE\x10\x1f\x12\x0b\n\x07\x44T_LIST\x10 \x12\x0b\n\x07\x44T_DICT\x10!\x12\x0b\n\x07\x44T_NONE\x10\"\x12\x0f\n\x0b\x44T_SYM_INST\x10#\x12\x0f\n\x0b\x44T_BASE_INT\x10$\x12\x10\n\x0c\x44T_BASE_UINT\x10%\x12\x11\n\rDT_BASE_FLOAT\x10&\x12\x0b\n\x07\x44T_TYPE\x10\'\x12\n\n\x06\x44T_ANY\x10(\x12\r\n\tDT_REFKEY\x10)\x12\n\n\x06\x44T_REF\x10*\x12\x10\n\x0c\x44T_COMPLEX64\x10+\x12\x11\n\rDT_COMPLEX128\x10,\x12\x13\n\x0f\x44T_BASE_COMPLEX\x10-\x12\x0f\n\x0b\x44T_BFLOAT16\x10.\x12\x10\n\x0c\x44T_BFLOATS16\x10/'
 )
 
 _VERSION = _descriptor.EnumDescriptor(
@@ -286,11 +286,21 @@ _DATATYPE = _descriptor.EnumDescriptor(
       serialized_options=None,
       type=None,
       create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='DT_BFLOAT16', index=46, number=46,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='DT_BFLOATS16', index=47, number=47,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
   ],
   containing_type=None,
   serialized_options=None,
   serialized_start=2650,
-  serialized_end=3365,
+  serialized_end=3400,
 )
 _sym_db.RegisterEnumDescriptor(_DATATYPE)
 
@@ -343,6 +353,8 @@ DT_REF = 42
 DT_COMPLEX64 = 43
 DT_COMPLEX128 = 44
 DT_BASE_COMPLEX = 45
+DT_BFLOAT16 = 46
+DT_BFLOATS16 = 47
 
 
 _INPUTPROTO_EDGETYPE = _descriptor.EnumDescriptor(

mindspore/train/callback/_backup_and_restore.py

@@ -50,25 +50,15 @@ class BackupAndRestore(Callback):
         >>> from mindspore.train import Model, BackupAndRestore, RunContext
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
         >>> model = Model(net, loss_fn=loss, optimizer=optim)
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> backup_ckpt = BackupAndRestore("backup")
-        >>> cb_params = {}
-        >>> cb_params["cur_epoch_num"] = 4
-        >>> cb_params["epoch_num"] = 4
-        >>> cb_params["cur_step_num"] = 2
-        >>> cb_params["batch_num"] = 2
-        >>> cb_params["net_outputs"] = Tensor(2.0)
-        >>> run_context = RunContext(cb_params)
-        >>> backup_ckpt.on_train_begin(run_context)
-        >>> backup_ckpt.on_train_epoch_end(run_context)
-        >>> backup_ckpt.on_train_end(run_context)
         >>> model.train(10, dataset, callbacks=backup_ckpt)
     """
     def __init__(self, backup_dir, save_freq="epoch", delete_checkpoint=True):

mindspore/train/callback/_callback.py

@@ -123,7 +123,8 @@ class Callback:
     recording current attributes. Users can add custimized attributes to the information.
     Training process can also be stopped by calling `request_stop` method. For details
     of custom Callback, please check
-    `Callback Mechanism <https://www.mindspore.cn/tutorials/en/r2.1/advanced/model/callback.html>`_.
+    `Callback tutorial <https://www.mindspore.cn/tutorials/en/r2.2/advanced/model/
+    callback.html#customized-callback-mechanism>`_.
 
     Examples:
         >>> import numpy as np
@@ -159,6 +160,10 @@ class Callback:
         Called once before the network executing.
         A backwards compatibility alias for `on_train_begin` and `on_eval_begin`.
 
+        Note:
+            `begin` is deprecated and will be deleted in a future version,
+            please use `on_train_begin` and `on_eval_begin` instead.
+
         Args:
             run_context (RunContext): Include some information of the model.
         """
@@ -168,6 +173,10 @@ class Callback:
         Called before each epoch beginning.
         A backwards compatibility alias for `on_train_epoch_begin` and `on_eval_epoch_begin`.
 
+        Note:
+            `epoch_begin` is deprecated and will be deleted in a future version,
+            please use `on_train_epoch_begin` and `on_eval_epoch_begin` instead.
+
         Args:
             run_context (RunContext): Include some information of the model.
         """
@@ -177,6 +186,10 @@ class Callback:
         Called after each epoch finished.
         A backwards compatibility alias for `on_train_epoch_end` and `on_eval_epoch_end`.
 
+        Note:
+            `epoch_end` is deprecated and will be deleted in a future version,
+            please use `on_train_epoch_end` and `on_eval_epoch_end` instead.
+
         Args:
             run_context (RunContext): Include some information of the model.
         """
@@ -186,6 +199,10 @@ class Callback:
         Called before each step beginning.
         A backwards compatibility alias for `on_train_step_begin` and `on_eval_step_begin`.
 
+        Note:
+            `step_begin` is deprecated and will be deleted in a future version,
+            please use `on_train_step_begin` and `on_eval_step_begin` instead.
+
         Args:
             run_context (RunContext): Include some information of the model.
         """
@@ -195,6 +212,10 @@ class Callback:
         Called after each step finished.
         A backwards compatibility alias for `on_train_step_end` and `on_eval_step_end`.
 
+        Note:
+            `step_end` is deprecated and will be deleted in a future version,
+            please use `on_train_step_end` and `on_eval_step_end` instead.
+
         Args:
             run_context (RunContext): Include some information of the model.
         """
@@ -204,6 +225,10 @@ class Callback:
         Called once after network training.
         A backwards compatibility alias for `on_train_end` and `on_eval_end`.
 
+        Note:
+            `end` is deprecated and will be deleted in a future version,
+            please use `on_train_end` and `on_eval_end` instead.
+
         Args:
             run_context (RunContext): Include some information of the model.
         """
@@ -468,7 +493,7 @@ class RunContext:
     `RunContext.original_args()` and add extra attributes to the information, but also can stop the
     training process by calling `request_stop` method. For details of custom Callback,
     please check
-    `Callback Mechanism <https://www.mindspore.cn/tutorials/en/r2.1/advanced/model/callback.html>`_.
+    `Callback Mechanism <https://www.mindspore.cn/tutorials/en/r2.2/advanced/model/callback.html>`_.
 
     `RunContext.original_args()` holds the model context information as a dictionary variable, and
     different attributes of the dictionary are stored in training or eval process. Details are as follows:
@@ -550,7 +575,7 @@ class RunContext:
 
         Tutorial Examples:
             - `Callback Mechanism - Customized Callback Mechanism
-              <https://mindspore.cn/tutorials/en/r2.1/advanced/model/callback.html#customized-callback-mechanism>`_
+              <https://mindspore.cn/tutorials/en/r2.2/advanced/model/callback.html#customized-callback-mechanism>`_
         """
         return self._original_args
 
@@ -563,7 +588,7 @@ class RunContext:
 
         Tutorial Examples:
             - `Callback Mechanism - Customized Training Termination Time
-              <https://mindspore.cn/tutorials/en/r2.1/advanced/model/callback.html#
+              <https://mindspore.cn/tutorials/en/r2.2/advanced/model/callback.html#
               customized-training-termination-time>`_
         """
         self._stop_requested = True

mindspore/train/callback/_checkpoint.py

@@ -62,7 +62,7 @@ def _chg_ckpt_file_name_if_same_exist(directory, prefix, exception=False):
                     suffix_num = max(suffix_num, int(num)+1)
 
     if suffix_num != 0:
-        prefix = prefix + "_" + str(suffix_num)
+        prefix = f'{prefix}_{suffix_num}'
 
     return prefix
 
@@ -94,8 +94,8 @@ class CheckpointConfig:
         async_save (bool): Whether asynchronous execution saves the checkpoint to a file. Default: ``False`` .
         saved_network (Cell): Network to be saved in checkpoint file. If the saved_network has no relation
             with the network in training, the initial value of saved_network will be saved. Default: ``None`` .
-        append_info (list): The information save to checkpoint file. Support "epoch_num", "step_num" and dict. The key
-            of dict must be str, the value of dict must be one of int, float, bool, Parameter or Tensor.
+        append_info (list): The information save to checkpoint file. Support "epoch_num", "step_num" and
+            dict. The key of dict must be str, the value of dict must be one of int, float, bool, Parameter or Tensor.
             Default: ``None`` .
         enc_key (Union[None, bytes]): Byte type key used for encryption. If the value is None, the encryption
                                       is not required. Default: ``None`` .
@@ -104,8 +104,6 @@ class CheckpointConfig:
         exception_save (bool): Whether to save the current checkpoint when an exception occurs. Default: ``False`` .
         kwargs (dict): Configuration options dictionary.
 
-            - incremental (bool): whether export checkpoint for MapParameter incrementally.
-
     Raises:
         ValueError: If input parameter is not the correct type.
 
@@ -114,13 +112,13 @@ class CheckpointConfig:
         >>> from mindspore.train import Model, CheckpointConfig, ModelCheckpoint
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
         >>> model = Model(net, loss_fn=loss, optimizer=optim)
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> config = CheckpointConfig(save_checkpoint_seconds=100, keep_checkpoint_per_n_minutes=5, saved_network=net)
         >>> config.save_checkpoint_steps
@@ -387,6 +385,21 @@ class ModelCheckpoint(Callback):
         ValueError: If `prefix` is not str or contains the '/' character.
         ValueError: If `directory` is not str.
         TypeError: If the config is not CheckpointConfig type.
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore.dataset as ds
+        >>> from mindspore import nn
+        >>> from mindspore.train import Model, ModelCheckpoint
+        >>>
+        >>> data = {"x": np.float32(np.random.rand(64, 10)), "y": np.random.randint(0, 5, (64,))}
+        >>> train_dataset = ds.NumpySlicesDataset(data=data).batch(32)
+        >>> net = nn.Dense(10, 5)
+        >>> crit = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
+        >>> opt = nn.Momentum(net.trainable_params(), 0.01, 0.9)
+        >>> ckpt_callback = ModelCheckpoint(prefix="myckpt")
+        >>> model = Model(network=net, optimizer=opt, loss_fn=crit)
+        >>> model.train(2, train_dataset, callbacks=[ckpt_callback])
     """
 
     def __init__(self, prefix='CKP', directory=None, config=None):
@@ -528,7 +541,9 @@ class ModelCheckpoint(Callback):
             self._last_time_for_keep = time.time()
             self._last_triggered_step = cb_params.cur_step_num
 
-            if context.get_context("enable_ge"):
+            # TODO(MS_DISABLE_REF_MODE): Delete when remove MS_DISABLE_REF_MODE env.
+            if context.get_context("enable_ge") and os.getenv('MS_DISABLE_REF_MODE') \
+                    and context.get_context("mode") == context.GRAPH_MODE:
                 set_cur_net(cb_params.train_network)
                 cb_params.train_network.exec_checkpoint_graph()
             if "epoch_num" in self._append_dict:

mindspore/train/callback/_early_stop.py

@@ -85,13 +85,13 @@ class EarlyStopping(Callback):
         >>> from mindspore import nn
         >>> from mindspore.train import Model, EarlyStopping
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
         >>> model = Model(net, loss_fn=loss, optimizer=optim, metrics={"acc"})
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> cb = EarlyStopping(monitor="acc", patience=3, verbose=True)
         >>> model.fit(10, dataset, callbacks=cb)

mindspore/train/callback/_landscape.py

@@ -187,10 +187,10 @@ class SummaryLandscape:
         ...     # If the device_target is Ascend, set the device_target to "Ascend"
         ...     ms.set_context(mode=ms.GRAPH_MODE, device_target="GPU")
         ...     # Create the dataset taking MNIST as an example. Refer to
-        ...     # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        ...     # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         ...     ds_train = create_dataset()
         ...     # Define the network structure of LeNet5. Refer to
-        ...     # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        ...     # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         ...     network = LeNet5()
         ...     net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
         ...     net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9)
@@ -210,13 +210,13 @@ class SummaryLandscape:
         ...     # Simple usage for visualization landscape:
         ...     def callback_fn():
         ...         # Define the network structure of LeNet5. Refer to
-        ...         # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        ...         # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         ...         network = LeNet5()
         ...         net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
         ...         metrics = {"Loss": Loss()}
         ...         model = Model(network, net_loss, metrics=metrics)
         ...         # Create the dataset taking MNIST as an example. Refer to
-        ...         # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        ...         # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         ...         ds_eval = create_dataset()
         ...         return model, network, ds_eval, metrics
         ...

mindspore/train/callback/_loss_monitor.py

@@ -43,13 +43,13 @@ class LossMonitor(Callback):
         >>> from mindspore.train import Model, LossMonitor
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
         >>> model = Model(net, loss_fn=loss, optimizer=optim)
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> loss_monitor = LossMonitor()
         >>> model.train(10, dataset, callbacks=loss_monitor)

mindspore/train/callback/_on_request_exit.py

@@ -55,13 +55,13 @@ class OnRequestExit(Callback):
         >>> import mindspore as ms
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
         >>> model = Model(net, loss_fn=loss, optimizer=optim)
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> on_request_exit = ms.train.OnRequestExit(file_name='LeNet5')
         >>> model.train(10, dataset, callbacks=on_request_exit)

mindspore/train/callback/_reduce_lr_on_plateau.py

@@ -54,7 +54,6 @@ class ReduceLROnPlateau(Callback):
             otherwise the valid monitor is ``"loss"``.
             When `monitor` is ``"loss"``, if train network has multiple outputs,
             the first element will be returned as training loss. Default: ``'eval_loss'``.
-
         factor (float): factor by which the learning rate will be reduced.
             `new_lr = lr * factor`. Default: ``0.1`` .
         patience (int): `monitor` value is better than history best value over
@@ -77,7 +76,7 @@ class ReduceLROnPlateau(Callback):
         min_lr (float): lower bound on the learning rate. Default: ``0`` .
 
     Raises:
-        ValueError: `mode` not in 'auto', 'min' or 'max'.
+        ValueError: `mode` not in ``'auto'``, ``'min'`` or ``'max'``.
         ValueError: The monitor value is not a scalar.
         ValueError: The learning rate is not a Parameter.
 
@@ -85,13 +84,13 @@ class ReduceLROnPlateau(Callback):
         >>> from mindspore import nn
         >>> from mindspore.train import Model, ReduceLROnPlateau
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
         >>> model = Model(net, loss_fn=loss, optimizer=optim, metrics={"acc"})
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> cb = ReduceLROnPlateau(monitor="acc", patience=3, verbose=True)
         >>> model.fit(10, dataset, callbacks=cb)

mindspore/train/callback/_summary_collector.py

@@ -71,9 +71,9 @@ class LineageMetadata:
 
 class SummaryCollector(Callback):
     """
-    SummaryCollector can help you to collect some common information.
+    SummaryCollector can help you to collect some common information,
+    such as loss, learning late, computational graph and so on.
 
-    It can help you to collect loss, learning late, computational graph and so on.
     SummaryCollector also enables the summary operator to collect data to summary files.
 
     Note:
@@ -92,6 +92,8 @@ class SummaryCollector(Callback):
             when (current steps % freq) equals to 0, and the first step will be collected at any time.
             It is important to note that if the data sink mode is used, the unit will become the `epoch`.
             It is not recommended to collect data too frequently, which can affect performance. Default: ``10`` .
+        num_process (int): Number of processes saving summary data. The more processes there are, the better the
+            performance, but there may be host memory overflow issues. Default: ``32`` .
         collect_specified_data (Union[None, dict]): Perform custom operations on the collected data.
             By default, if set to None, all data is collected as the default behavior.
             You can customize the collected data with a dictionary.
@@ -104,10 +106,12 @@ 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/r2.1/lineage_and_scalars_comparison.html>`_
+              <https://www.mindspore.cn/mindinsight/docs/en/r2.2/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`` .
+              this field will be displayed on the `lineage page
+              <https://www.mindspore.cn/mindinsight/docs/en/r2.2/lineage_and_scalars_comparison.html>`_
+              of MindInsight. Default: ``True`` .
             - collect_input_data (bool): Whether to collect dataset for each training.
               Currently only image data is supported.
               If there are multiple columns of data in the dataset, the first column should be image data.
@@ -147,7 +151,8 @@ class SummaryCollector(Callback):
             False: it means that after specified data is set, only the specified data is collected,
             and the others are not collected. Default: ``True`` .
         custom_lineage_data (Union[dict, None]): Allows you to customize the data and present it on the MingInsight
-            lineage page. In the custom data, the type of the key supports str, and the type of value supports str, int
+            `lineage page <https://www.mindspore.cn/mindinsight/docs/en/r2.2/lineage_and_scalars_comparison.html>`_ .
+            In the custom data, the type of the key supports str, and the type of value supports str, int
             and float. Default: ``None`` , it means there is no custom data.
         collect_tensor_freq (Optional[int]): The same semantics as the `collect_freq`, but controls TensorSummary only.
             Because TensorSummary data is too large to be compared with other summary data, this parameter is used to
@@ -186,10 +191,10 @@ class SummaryCollector(Callback):
         ...     ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend")
         ...     mnist_dataset_dir = '/path/to/mnist_dataset_directory'
         ...     # Create the dataset taking MNIST as an example. Refer to
-        ...     # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        ...     # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         ...     ds_train = create_dataset()
         ...     # Define the network structure of LeNet5. Refer to
-        ...     # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        ...     # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         ...     network = LeNet5(10)
         ...     net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
         ...     net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9)
@@ -225,6 +230,7 @@ class SummaryCollector(Callback):
     def __init__(self,
                  summary_dir,
                  collect_freq=10,
+                 num_process=32,
                  collect_specified_data=None,
                  keep_default_action=True,
                  custom_lineage_data=None,
@@ -278,6 +284,7 @@ class SummaryCollector(Callback):
         self._is_parse_loss_success = True
         self._first_step = True
         self._dataset_sink_mode = True
+        self._num_process = num_process
 
     def __enter__(self):
         self._record = SummaryRecord(log_dir=self._summary_dir,
@@ -1150,7 +1157,7 @@ class SummaryCollector(Callback):
         except TypeError as exc:
             logger.warning("Summary cannot collect the type of metrics, currently support type: dict, list, tuple, "
                            "str, int, float, bool and None. %s.", str(exc))
-        self._parse_dataset(cb_params, eval_lineage)
+        _ = self._parse_dataset(cb_params, eval_lineage)
 
         eval_lineage_message = self._package_eval_lineage_message(eval_lineage)
         self._record.add_value(PluginEnum.EVAL_LINEAGE.value, 'eval_lineage', eval_lineage_message)

mindspore/train/callback/_time_monitor.py

@@ -30,30 +30,63 @@ class TimeMonitor(Callback):
             if the program get `batch_num` during training, `data_size` will be set to `batch_num`,
             otherwise `data_size` will be used. Default: ``None`` .
 
+        data_time (bool): Whether to sow the average time of fetching data in Host.
+            Note that data fetch and network compute are processed sequentially in non dataset sink mode, while
+            they are asynchronous in dataset sink mode. Default: ``False`` .
+
     Raises:
         ValueError: If data_size is not positive int.
+        TypeError: If data_time is not bool.
 
     Examples:
         >>> from mindspore import nn
         >>> from mindspore.train import Model, TimeMonitor
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
         >>> model = Model(net, loss_fn=loss, optimizer=optim)
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> time_monitor = TimeMonitor()
         >>> model.train(10, dataset, callbacks=time_monitor)
     """
 
-    def __init__(self, data_size=None):
+    def __init__(self, data_size=None, data_time=False):
         super(TimeMonitor, self).__init__()
         self.data_size = data_size
         self.epoch_time = time.time()
+        self.data_time = data_time
+        self.data_time_sum = 0.0
+        self.data_time_start = 0.0
+        self.data_sink = lambda c: c.original_args()["dataset_sink_mode"]
+        Validator.check_bool(data_time, "data_time")
+
+    def on_train_step_begin(self, run_context):
+        """
+        Record time at the beginning of step.
+
+        Args:
+            run_context (RunContext): Context of the process running. For more details,
+                    please refer to :class:`mindspore.train.RunContext`.
+        """
+        if self.data_time and not self.data_sink(run_context):
+            interval = time.time() - self.data_time_start
+            self.data_time_sum = self.data_time_sum + interval
+
+    def on_train_step_end(self, run_context):
+        """
+        Record time at the end of step.
+
+        Args:
+            run_context (RunContext): Context of the process running. For more details,
+                    please refer to :class:`mindspore.train.RunContext`.
+        """
+        if self.data_time and not self.data_sink(run_context):
+            self.data_time_start = time.time()
 
     def epoch_begin(self, run_context):
         """
@@ -64,6 +97,9 @@ class TimeMonitor(Callback):
                     please refer to :class:`mindspore.train.RunContext`.
         """
         self.epoch_time = time.time()
+        if self.data_time and not self.data_sink(run_context):
+            self.data_time_sum = 0.0
+            self.data_time_start = time.time()
 
     def epoch_end(self, run_context):
         """
@@ -84,5 +120,22 @@ class TimeMonitor(Callback):
         Validator.check_positive_int(step_size)
 
         step_seconds = epoch_seconds / step_size
-        print("{} epoch time: {:5.3f} ms, per step time: {:5.3f} ms".format
-              (mode.title(), epoch_seconds, step_seconds), flush=True)
+
+        train_log = "{} epoch time: {:5.3f} ms, per step time: {:5.3f} ms".format(
+            mode.title(), epoch_seconds, step_seconds)
+
+        if self.data_time and not self.data_sink(run_context):
+            data_step_seconds = self.data_time_sum * 1000 / step_size
+            data_log = " (data time: {:5.3f} ms)".format(data_step_seconds)
+            train_log += data_log
+        elif self.data_time and self.data_sink(run_context):
+            # send info viewer to query epoch message of cur_epoch_num
+            send_info = cb_params["dataset_helper"].get_send_info(run_context)
+            epoch = cb_params["cur_epoch_num"]
+            epoch_send_info = send_info.epoch(epoch)
+            # show average time of fetching data time
+            fetch_data_time = epoch_send_info["fetch_data_time"]
+            data_log = " (data time: {:5.3f} ms)".format(fetch_data_time)
+            train_log += data_log
+
+        print(train_log, flush=True)

mindspore/train/data_sink.py

@@ -130,16 +130,10 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
     A wrapper function to generate a function for the input function.
 
     Note:
-        When using data sinking, the dataset will automatically cycle. At this time,
-        only the total number of training steps (total_step) and the number of steps for each sinking (sink_size)
-        need to be considered. When switching from training by rounds (epochs) to data sinking,
-        the formula is as follows:
-
-        total_step = epochs * dataset_size
-
-        train_sink_step = total_step / sink_size
-
-        After transforming from `mindspore.data_sink`, you need to execute `train_sink_step` step for training.
+        When using data sinking, the dataset will be automatically sent in a loop, and only the step size of sinking
+        `sink_size` needs to be considered. The default value of `sink_size` is ``1``, which means that all data will
+        be sunk every epoch. If `sink_size` is greater than 1, the amount of data sunk per epoch will be the dataset
+        with a size of `sink_size`.
 
     Args:
         fn (Function): The Python function that will be run with dataset.
@@ -212,7 +206,7 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
         real_sink_fun = _get_sink_fun(sink_fun, key_info, is_info_queue, dataset, jit_config)
 
         loop = sink_size
-        if jit_config is not None:
+        if jit_config is not None and context.get_context('mode') == context.GRAPH_MODE:
             loop = 1
 
         out = None