mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__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*\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/'
+  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*\x89\x06\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/\x12\x0b\n\x07\x44T_INT4\x10\x30\x12\x0c\n\x08\x44T_SLICE\x10\x31'
 )
 
 _VERSION = _descriptor.EnumDescriptor(
@@ -296,11 +296,21 @@ _DATATYPE = _descriptor.EnumDescriptor(
       serialized_options=None,
       type=None,
       create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT4', index=48, number=48,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='DT_SLICE', index=49, number=49,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
   ],
   containing_type=None,
   serialized_options=None,
   serialized_start=2650,
-  serialized_end=3400,
+  serialized_end=3427,
 )
 _sym_db.RegisterEnumDescriptor(_DATATYPE)
 
@@ -355,6 +365,8 @@ DT_COMPLEX128 = 44
 DT_BASE_COMPLEX = 45
 DT_BFLOAT16 = 46
 DT_BFLOATS16 = 47
+DT_INT4 = 48
+DT_SLICE = 49
 
 
 _INPUTPROTO_EDGETYPE = _descriptor.EnumDescriptor(

mindspore/train/callback/__init__.py

@@ -35,7 +35,10 @@ from mindspore.train.callback._early_stop import EarlyStopping
 from mindspore.train.callback._reduce_lr_on_plateau import ReduceLROnPlateau
 from mindspore.train.callback._on_request_exit import OnRequestExit
 from mindspore.train.callback._backup_and_restore import BackupAndRestore
+from mindspore.train.callback._flops_collector import FlopsUtilizationCollector
+from mindspore.train.callback._tft_register import TFTRegister
 
-__all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint",
+__all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint", "FlopsUtilizationCollector",
            "SummaryCollector", "CheckpointConfig", "RunContext", "LearningRateScheduler", "SummaryLandscape",
-           "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping", "OnRequestExit", "BackupAndRestore"]
+           "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping", "OnRequestExit", "BackupAndRestore",
+           "TFTRegister"]

mindspore/train/callback/_backup_and_restore.py

@@ -34,15 +34,15 @@ class BackupAndRestore(Callback):
 
     Args:
         backup_dir (str): Path to store and load the checkpoint file.
-        save_freq(Union['epoch', int]): When set to 'epoch' the callback saves the checkpoint at the end of
+        save_freq (Union["epoch", int]): When set to ``"epoch"`` the callback saves the checkpoint at the end of
                                         each epoch. When set to an integer, the callback saves the checkpoint
                                         every `save_freq` epoch. Default: ``"epoch"`` .
-        delete_checkpoint(bool): If `delete_checkpoint=True`, the checkpoint will be deleted after
+        delete_checkpoint (bool): If `delete_checkpoint=True`, the checkpoint will be deleted after
                                         training is finished. Default: ``True`` .
 
     Raises:
         ValueError: If backup_dir is not str.
-        ValueError: If save_freq is not 'epoch' or int.
+        ValueError: If save_freq is not ``"epoch"`` or int.
         ValueError: If delete_checkpoint is not bool.
 
     Examples:
@@ -50,13 +50,13 @@ 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.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/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.2/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> backup_ckpt = BackupAndRestore("backup")
         >>> model.train(10, dataset, callbacks=backup_ckpt)

mindspore/train/callback/_callback.py

@@ -123,7 +123,7 @@ 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 tutorial <https://www.mindspore.cn/tutorials/en/r2.2/advanced/model/
+    `Callback tutorial <https://www.mindspore.cn/docs/en/master/model_train/train_process/model/
     callback.html#customized-callback-mechanism>`_.
 
     Examples:
@@ -493,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.2/advanced/model/callback.html>`_.
+    `Callback Mechanism <https://www.mindspore.cn/docs/en/master/model_train/train_process/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:
@@ -575,7 +575,7 @@ class RunContext:
 
         Tutorial Examples:
             - `Callback Mechanism - Customized Callback Mechanism
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/model/callback.html#customized-callback-mechanism>`_
+              <https://mindspore.cn/docs/en/master/model_train/train_process/model/callback.html#customized-callback-mechanism>`_
         """
         return self._original_args
 
@@ -588,7 +588,7 @@ class RunContext:
 
         Tutorial Examples:
             - `Callback Mechanism - Customized Training Termination Time
-              <https://mindspore.cn/tutorials/en/r2.2/advanced/model/callback.html#
+              <https://mindspore.cn/docs/en/master/model_train/train_process/model/callback.html#
               customized-training-termination-time>`_
         """
         self._stop_requested = True

mindspore/train/callback/_checkpoint.py

@@ -18,8 +18,8 @@ from __future__ import absolute_import
 import os
 import stat
 import time
-
 import threading
+
 import mindspore.context as context
 from mindspore import log as logger
 from mindspore import nn
@@ -28,19 +28,49 @@ from mindspore.train._utils import _make_directory
 from mindspore.train.serialization import save_checkpoint, _save_graph
 from mindspore.parallel._cell_wrapper import destroy_allgather_cell
 from mindspore.parallel._recovery_context import _set_recovery_context, _get_recovery_context
+from mindspore.parallel._auto_parallel_context import _get_auto_parallel_context
+from mindspore.parallel._utils import _get_device_num
+from mindspore.communication.management import get_rank
+from mindspore.train._utils import get_parameter_redundancy, remove_param_redundancy
 from mindspore.train.callback._callback import Callback, set_cur_net
 from mindspore.common.tensor import Tensor
 from mindspore.common.parameter import Parameter
-from mindspore._c_expression import _collect_host_info
-
+from mindspore.common.generator import Generator
+from mindspore.common.api import _cell_graph_executor
+from mindspore._c_expression import collect_host_info, get_clock_syscnt
 
 _cur_dir = os.getcwd()
 SAVE_DIR = _cur_dir
 _info_list = ["epoch_num", "step_num"]
 
 
+def _get_dp_tp_from_redundancy(redundancy_tuple):
+    """From redundancy get dp and tp"""
+    dp = []
+    tp = []
+    for dp_value in redundancy_tuple:
+        dp.append(list(dp_value))
+    for i in range(len(redundancy_tuple[0])):
+        tp.append([v[i] for v in redundancy_tuple])
+    return dp, tp
+
+
+def _get_dp_tp_from_layout(parameter_redundancy_dict):
+    """From layout dict get dp and tp"""
+    tp = []
+    dp = []
+    value_len = 0
+    for _, value in parameter_redundancy_dict.items():
+        if len(value) > value_len:
+            value_len = len(value)
+            dp, tp = _get_dp_tp_from_redundancy(value)
+    return dp, tp
+
+
 def _chg_ckpt_file_name_if_same_exist(directory, prefix, exception=False):
     """Check if there is a file with the same name."""
+    if callable(prefix) or callable(directory):
+        return prefix
     files = os.listdir(directory)
     suffix_num = 0
     pre_len = len(prefix)
@@ -57,9 +87,9 @@ def _chg_ckpt_file_name_if_same_exist(directory, prefix, exception=False):
             if index == 0:
                 suffix_num = max(suffix_num, 1)
             elif index != -1:
-                num = filename[pre_len+1:pre_len+index]
+                num = filename[pre_len + 1:pre_len + index]
                 if num.isdigit():
-                    suffix_num = max(suffix_num, int(num)+1)
+                    suffix_num = max(suffix_num, int(num) + 1)
 
     if suffix_num != 0:
         prefix = f'{prefix}_{suffix_num}'
@@ -67,19 +97,28 @@ def _chg_ckpt_file_name_if_same_exist(directory, prefix, exception=False):
     return prefix
 
 
+def _check_format_and_other_params(format, enc_key, enc_mode, crc_check=False, async_save=False, exception_save=False,
+                                   map_param_inc=False, global_step_num=None):
+    param_not_default = (enc_key is not None or enc_mode != "AES-GCM" or crc_check or async_save
+                         or exception_save or map_param_inc or global_step_num is not None)
+    if format == "safetensors" and param_not_default:
+        raise ValueError("For 'save_checkpoint', when format is 'safetensors', other param must be default.")
+
+
 class CheckpointConfig:
     """
     The configuration of model checkpoint.
 
     Note:
-        During the training process, if dataset is transmitted through the data channel,
-        it is suggested to set 'save_checkpoint_steps' to an integer multiple of loop_size.
-        Otherwise, the time to save the checkpoint may be biased.
-        It is recommended to set only one save strategy and one keep strategy at the same time.
-        If both `save_checkpoint_steps` and `save_checkpoint_seconds` are set,
-        `save_checkpoint_seconds` will be invalid.
-        If both `keep_checkpoint_max` and `keep_checkpoint_per_n_minutes` are set,
-        `keep_checkpoint_per_n_minutes` will be invalid.
+        - During the training process, if dataset is transmitted through the data channel,
+          it is suggested to set 'save_checkpoint_steps' to an integer multiple of loop_size.
+          Otherwise, the time to save the checkpoint may be biased.
+          It is recommended to set only one save strategy and one keep strategy at the same time.
+          If both `save_checkpoint_steps` and `save_checkpoint_seconds` are set,
+          `save_checkpoint_seconds` will be invalid.
+          If both `keep_checkpoint_max` and `keep_checkpoint_per_n_minutes` are set,
+          `keep_checkpoint_per_n_minutes` will be invalid.
+        - The `enc_mode` and `crc_check` parameters are mutually exclusive and cannot be configured simultaneously.
 
     Args:
         save_checkpoint_steps (int): Steps to save checkpoint. Default: ``1`` .
@@ -102,6 +141,12 @@ class CheckpointConfig:
         enc_mode (str): This parameter is valid only when enc_key is not set to None. Specifies the encryption
                         mode, currently supports 'AES-GCM', 'AES-CBC' and 'SM4-CBC'. Default: ``'AES-GCM'`` .
         exception_save (bool): Whether to save the current checkpoint when an exception occurs. Default: ``False`` .
+        crc_check (bool): Whether to perform crc32 calculation when saving checkpoint and save the calculation
+                          result to the end of ckpt. Default: ``False`` .
+        remove_redundancy (bool): Whether to enable saving the checkpoint with redundancy removal.
+            Redundancy removal refers to eliminating redundant data in data parallelism mode. Default: ``False`` , means
+            redundant-free saving is not enabled.
+        format (str): Format of the output file, can be "ckpt" or "safetensors". Default: "ckpt".
         kwargs (dict): Configuration options dictionary.
 
     Raises:
@@ -112,23 +157,21 @@ class CheckpointConfig:
         >>> from mindspore.train import Model, CheckpointConfig, ModelCheckpoint
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/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.2/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/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
         1
         >>> config.save_checkpoint_seconds
-        100
         >>> config.keep_checkpoint_max
         5
         >>> config.keep_checkpoint_per_n_minutes
-        5
         >>> config.integrated_save
         True
         >>> config.async_save
@@ -155,6 +198,9 @@ class CheckpointConfig:
                  enc_key=None,
                  enc_mode='AES-GCM',
                  exception_save=False,
+                 crc_check=False,
+                 remove_redundancy=False,
+                 format="ckpt",
                  **kwargs):
 
         if save_checkpoint_steps is not None:
@@ -197,7 +243,14 @@ class CheckpointConfig:
         self._append_dict = self._handle_append_info(append_info)
         self._enc_key = Validator.check_isinstance('enc_key', enc_key, (type(None), bytes))
         self._enc_mode = Validator.check_isinstance('enc_mode', enc_mode, str)
+        self._crc_check = Validator.check_isinstance('crc_check', crc_check, bool)
+        self._format = Validator.check_isinstance('format', format, str)
         self._map_param_inc = kwargs.get('incremental', False)
+        self.enable_redundance = kwargs.get('enable_redundance', False)
+        self.remove_redundancy = Validator.check_isinstance('remove_redundancy', remove_redundancy, bool)
+
+        _check_format_and_other_params(format, enc_key, enc_mode, crc_check, async_save, exception_save,
+                                       self._map_param_inc)
 
     @property
     def save_checkpoint_steps(self):
@@ -288,6 +341,20 @@ class CheckpointConfig:
         """
         return self._enc_mode
 
+    @property
+    def crc_check(self):
+        """
+        Get the value of the whether to enable crc check.
+
+        Returns:
+            bool, whether to enable crc check.
+        """
+        return self._crc_check
+
+    @property
+    def format(self):
+        return self._format
+
     @property
     def append_dict(self):
         """
@@ -353,7 +420,8 @@ class CheckpointConfig:
                     raise TypeError(f"For 'CheckpointConfig', the element of 'append_info' must has only one dict, "
                                     "but got {dict_num}")
                 for key, value in element.items():
-                    if isinstance(key, str) and isinstance(value, (int, float, bool, str, Parameter, Tensor)):
+                    if isinstance(key, str) and isinstance(value,
+                                                           (int, float, bool, str, Parameter, Tensor, Generator)):
                         handle_append_info[key] = value
                     else:
                         raise TypeError(f"For 'CheckpointConfig', the key type of the dict 'append_info' "
@@ -376,14 +444,16 @@ class ModelCheckpoint(Callback):
         parameters of the optimizer by default.
 
     Args:
-        prefix (str): The prefix name of checkpoint files. Default: ``'CKP'`` .
-        directory (str): The path of the folder which will be saved in the checkpoint file.
-            By default, the file is saved in the current directory. Default: ``None`` .
+        prefix (Union[str, callable object]): The prefix name or callable object to generate name of checkpoint files.
+            Default: ``'CKP'`` .
+        directory (Union[str, callable object]): The folder path where the checkpoint is stored, or the callable object
+            used to generate the path. By default, the file is saved in the current directory.
+            Default: ``None`` .
         config (CheckpointConfig): Checkpoint strategy configuration. Default: ``None`` .
 
     Raises:
-        ValueError: If `prefix` is not str or contains the '/' character.
-        ValueError: If `directory` is not str.
+        ValueError: If `prefix` is not str or contains the '/' character and is not a callable object.
+        ValueError: If `directory` is not str and is not a callable object.
         TypeError: If the config is not CheckpointConfig type.
 
     Examples:
@@ -409,19 +479,29 @@ class ModelCheckpoint(Callback):
         self._last_time = time.time()
         self._last_time_for_keep = time.time()
         self._last_triggered_step = 0
+        """a callable for users to set self-defined prefix."""
+        self._prefix_func = None
+        """a callable for users to set self-defined directory."""
+        self._directory_func = None
 
-        if not isinstance(prefix, str) or prefix.find('/') >= 0:
+        if not callable(prefix) and (not isinstance(prefix, str) or prefix.find('/') >= 0):
             raise ValueError("For 'ModelCheckpoint', the argument 'prefix' "
                              "for checkpoint file name is invalid, it must be "
-                             "string and does not contain '/', but got {}.".format(prefix))
+                             "callable or string that does not contain '/', but got {}.".format(prefix))
         self._prefix = prefix
         self._exception_prefix = prefix
 
         if directory is not None:
-            self._directory = _make_directory(directory)
+            if callable(directory):
+                self._directory_func = directory
+            else:
+                self._directory = _make_directory(directory)
         else:
             self._directory = _cur_dir
 
+        if callable(prefix):
+            self._prefix_func = prefix
+
         if _get_recovery_context("enable_recovery"):
             _set_recovery_context(ckpt_path=self._directory)
 
@@ -434,9 +514,15 @@ class ModelCheckpoint(Callback):
                                 "but got {}.".format(type(config)))
             self._config = config
 
+        self._aiturbo_init_flag = os.getenv("AITURBO") == "1"
         # get existing checkpoint files
-        self._manager = CheckpointManager()
-        self._prefix = _chg_ckpt_file_name_if_same_exist(self._directory, self._prefix)
+        if self._aiturbo_init_flag:
+            from aiturbo.checkpoint.aiturbo_mindspore_ckpt import CheckpointShmManager
+            self._manager = CheckpointShmManager()
+        else:
+            self._manager = CheckpointManager(self._config.format)
+        if not callable(directory) and not callable(prefix):
+            self._prefix = _chg_ckpt_file_name_if_same_exist(self._directory, self._prefix)
         self._append_dict = self._config.append_dict or {}
         self._append_epoch_num = self._append_dict.get("epoch_num") if "epoch_num" in self._append_dict else 0
         self._append_step_num = self._append_dict.get("step_num") if "step_num" in self._append_dict else 0
@@ -452,7 +538,37 @@ class ModelCheckpoint(Callback):
             run_context (RunContext): Context of the train running.
         """
         cb_params = run_context.original_args()
-        _collect_host_info("Callback", "ModelCheckpoint", "step_end", level=1)
+        if self._aiturbo_init_flag:
+            from aiturbo.checkpoint import aiturbo_mindspore as aiturbo
+            ckpt_storage_path = self._directory
+            rank_id = get_rank()
+            stage_num = _get_auto_parallel_context("pipeline_stages")
+            stage_rank_num = _get_device_num() // stage_num
+            param_layout = cb_params.train_network.parameter_layout_dict
+            if not param_layout:
+                layout = {"stage_num": stage_num, "stage_rank_num": stage_rank_num, "stage_layout": None}
+                aiturbo.init(ckpt_storage_path, rank_id, layout, None, False, None)
+            else:
+                device_num = _get_device_num()
+                chunk_size = device_num // stage_num
+                initial_rank = (rank_id // chunk_size) * chunk_size
+                param_redundancy_dict = get_parameter_redundancy(param_layout, initial_rank)
+                dp, _ = _get_dp_tp_from_layout(param_redundancy_dict)
+                layout = {"stage_num": stage_num, "stage_rank_num": stage_rank_num,
+                          "stage_layout": param_redundancy_dict}
+                single_params = remove_param_redundancy(param_redundancy_dict)
+                single_params = {device_id: list(params) for device_id, params in single_params.items()}
+                aiturbo.init(ckpt_storage_path, rank_id, layout, single_params, not self._config.enable_redundance, dp)
+            self._aiturbo_init_flag = False
+        if self._prefix_func:
+            self._prefix = self._prefix_func(cb_params)
+            if not isinstance(self._prefix, str) or self._prefix.find('/') >= 0:
+                raise ValueError("For 'ModelCheckpoint', the argument 'prefix' "
+                                 "for checkpoint file name is callable, it must return a "
+                                 "string that does not contain '/', but got {}.".format(self._prefix))
+        if self._directory_func:
+            self._directory = self._directory_func(cb_params)
+        collect_host_info("Callback", "ModelCheckpoint", "step_end", start_time=get_clock_syscnt(), level=1)
         # In disaster recovery scenario, the training process may be rolled back to the last step where
         # the ckpt was successfully saved, so the _last_triggered_step should be updated.
         if _get_recovery_context("enable_recovery") and cb_params.last_save_ckpt_step is not None:
@@ -481,7 +597,7 @@ class ModelCheckpoint(Callback):
             run_context (RunContext): Context of the train running.
         """
         cb_params = run_context.original_args()
-        _collect_host_info("Callback", "ModelCheckpoint", "end", level=1)
+        collect_host_info("Callback", "ModelCheckpoint", "end", start_time=get_clock_syscnt(), level=1)
         _to_save_last_ckpt = True
 
         self._save_ckpt(cb_params, _to_save_last_ckpt)
@@ -507,6 +623,13 @@ class ModelCheckpoint(Callback):
 
         return False
 
+    def _append_dict_content(self, epoch_num, step_num):
+        """Append append_dict content."""
+        if "epoch_num" in self._append_dict:
+            self._append_dict["epoch_num"] = self._append_epoch_num + epoch_num
+        if "step_num" in self._append_dict:
+            self._append_dict["step_num"] = self._append_step_num + step_num
+
     def _save_ckpt(self, cb_params, force_to_save=False):
         """Save checkpoint files."""
         if cb_params.cur_step_num == self._last_triggered_step:
@@ -520,8 +643,11 @@ class ModelCheckpoint(Callback):
         step_num_in_epoch = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
 
         if save_ckpt:
-            cur_ckpoint_file = self._prefix + "-" + str(cb_params.cur_epoch_num) + "_" \
-                + str(step_num_in_epoch) + ".ckpt"
+            if self._prefix_func:
+                cur_ckpoint_file = self._prefix + f".{self._config.format}"
+            else:
+                cur_ckpoint_file = self._prefix + "-" + str(cb_params.cur_epoch_num) + "_" \
+                                   + str(step_num_in_epoch) + f".{self._config.format}"
             # update checkpoint file list.
             self._manager.update_ckpoint_filelist(self._directory, self._prefix)
             # keep checkpoint files number equal max number.
@@ -545,15 +671,53 @@ class ModelCheckpoint(Callback):
             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:
-                self._append_dict["epoch_num"] = self._append_epoch_num + cb_params.cur_epoch_num
-            if "step_num" in self._append_dict:
-                self._append_dict["step_num"] = self._append_step_num + cb_params.cur_step_num
+                cb_params.train_network.add_flags(ge_sync_data=True)
+                _cell_graph_executor(cb_params.train_network, phase='save')
+            self._append_dict_content(cb_params.cur_epoch_num, cb_params.cur_step_num)
             network = self._config.saved_network if self._config.saved_network is not None else cb_params.train_network
-            save_checkpoint(network, cur_file, self._config.integrated_save, self._config.async_save,
-                            self._append_dict, self._config.enc_key, self._config.enc_mode,
-                            incremental=self._map_param_inc)
+            if os.getenv("AITURBO") == "1":
+                save_checkpoint(network, cur_file, self._config.integrated_save, self._config.async_save,
+                                self._append_dict, self._config.enc_key, self._config.enc_mode,
+                                crc_check=self._config.crc_check, incremental=self._map_param_inc,
+                                global_step_num=cb_params.cur_step_num)
+            elif self._config.remove_redundancy:
+                parallel_mode = context.get_auto_parallel_context("parallel_mode")
+                if parallel_mode == "stand_alone":
+                    raise TypeError(f"The deduplication feature for saving checkpoint can only be used "
+                                    f"in parallel scenarios, but got {parallel_mode}.")
+                param_layout = network.parameter_layout_dict
+                rank_id = get_rank()
+                if param_layout:
+                    device_num = _get_device_num()
+                    stage_num = _get_auto_parallel_context("pipeline_stages")
+                    chunk_size = device_num // stage_num
+                    initial_rank = (rank_id // chunk_size) * chunk_size
+                    param_redundancy_dict = get_parameter_redundancy(param_layout, initial_rank)
+                    single_params = remove_param_redundancy(param_redundancy_dict)
+                    save_param_names = single_params.get(rank_id)
+                    param_layout_set = set(param_layout.keys())
+                    if save_param_names == param_layout.keys():
+                        logger.warning(
+                            f"For remove_redundancy save checkpoint, the saved parameters are non-redundant.")
+
+                    def choice_func(x):
+                        return x not in param_layout_set or x in save_param_names
+                else:
+                    param_redundancy_dict = get_parameter_redundancy(network)
+                    single_params = remove_param_redundancy(param_redundancy_dict)
+                    save_param_names = single_params.get(rank_id)
+
+                    def choice_func(x):
+                        return x in save_param_names
+                save_checkpoint(network, cur_file, False, self._config.async_save,
+                                self._append_dict, self._config.enc_key, self._config.enc_mode,
+                                crc_check=self._config.crc_check, format=self._config.format,
+                                incremental=self._map_param_inc, choice_func=choice_func)
+            else:
+                save_checkpoint(network, cur_file, self._config.integrated_save, self._config.async_save,
+                                self._append_dict, self._config.enc_key, self._config.enc_mode,
+                                crc_check=self._config.crc_check, format=self._config.format,
+                                incremental=self._map_param_inc)
 
             self._latest_ckpt_file_name = cur_file
 
@@ -573,12 +737,23 @@ class ModelCheckpoint(Callback):
         """Return the latest checkpoint path and file name."""
         return self._latest_ckpt_file_name
 
+    @property
+    def _get_save_checkpoint_steps(self):
+        """Return save ckpt steps"""
+        return self._config.save_checkpoint_steps
+
+    @property
+    def _get_last_trigger_step(self):
+        """Return last triggered steps"""
+        return self._last_triggered_step
+
 
 class CheckpointManager:
     """Manage checkpoint files according to train_config of checkpoint."""
 
-    def __init__(self):
+    def __init__(self, format='ckpt'):
         self._ckpoint_filelist = []
+        self._format = format
 
     @property
     def ckpoint_filelist(self):
@@ -593,10 +768,12 @@ class CheckpointManager:
     def update_ckpoint_filelist(self, directory, prefix):
         """Update the checkpoint file list."""
         self._ckpoint_filelist = []
+        format = self._format
+        format_length = len(format) + 1
         files = os.listdir(directory)
         for filename in files:
-            if os.path.splitext(filename)[-1] == ".ckpt" and filename.startswith(prefix + "-"):
-                mid_name = filename[len(prefix):-5]
+            if os.path.splitext(filename)[-1] == f".{format}" and filename.startswith(prefix + "-"):
+                mid_name = filename[len(prefix):-format_length]
                 flag = not (True in [char.isalpha() for char in mid_name])
                 if flag:
                     self._ckpoint_filelist.append(os.path.join(directory, filename))