U Ph| @sddlmZddlZddlZddlZddlmZmZddlm Z m Z ddl Z ddl m Z ddlmZdZdd d d d d dddddZdd ddddddZdS)) annotationsN)MappingSequence)IOAny)get_config_values)JITMetadataKeysz metadata.jsonTFztorch.nn.Modulez str | IO[Any]boolzMapping[str, Any] | NonezMapping[str, bytes] | NoneNone)jit_objfilename_prefix_or_streaminclude_config_valsappend_timestamp meta_valuesmore_extra_filesreturnc Cstj}i}|r2|t||tjj<|dk rD||t |}t | i} |dk rl| |t |trtj|\} } | dkrd} |r|| d| }n| | }tj||| dS)a. Save the JIT object (script or trace produced object) `jit_obj` to the given file or stream with metadata included as a JSON file. The Torchscript format is a zip file which can contain extra file data which is used here as a mechanism for storing metadata about the network being saved. The data in `meta_values` should be compatible with conversion to JSON using the standard library function `dumps`. The intent is this metadata will include information about the network applicable to some use case, such as describing the input and output format, a network name and version, a plain language description of what the network does, and other relevant scientific information. Clients can use this information to determine automatically how to use the network, and users can read what the network does and keep track of versions. Examples:: net = torch.jit.script(monai.networks.nets.UNet(2, 1, 1, [8, 16], [2])) meta = { "name": "Test UNet", "used_for": "demonstration purposes", "input_dims": 2, "output_dims": 2 } # save the Torchscript bundle with the above dictionary stored as an extra file save_net_with_metadata(m, "test", meta_values=meta) # load the network back, `loaded_meta` has same data as `meta` plus version information loaded_net, loaded_meta, _ = load_net_with_metadata("test.ts") Args: jit_obj: object to save, should be generated by `script` or `trace`. filename_prefix_or_stream: filename or file-like stream object, if filename has no extension it becomes `.ts`. include_config_vals: if True, MONAI, Pytorch, and Numpy versions are included in metadata. append_timestamp: if True, a timestamp for "now" is appended to the file's name before the extension. meta_values: metadata values to store with the object, not limited just to keys in `JITMetadataKeys`. more_extra_files: other extra file data items to include in bundle, see `_extra_files` of `torch.jit.save`. Nz.tsz _%Y%m%d%H%M%S)datetimenowupdater astimezone isoformatr TIMESTAMPvaluejsondumpsMETADATA_FILENAMEencode isinstancestrospathsplitextstrftimetorchjitsave) r r r rrrrmetadict json_data extra_filesZfilename_no_extextr+Q/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/data/torchscript_utils.pysave_net_with_metadatas&-       r-r+ztorch.device | Nonez Sequence[str]z"tuple[torch.nn.Module, dict, dict])r map_locationrrcCsbdd|D}d|t<tj|||}t|}t|krJ|t}|t=nd}t|}|||fS)a] Load the module object from the given Torchscript filename or stream, and convert the stored JSON metadata back to a dict object. This will produce an empty dict if the metadata file is not present. Args: filename_prefix_or_stream: filename or file-like stream object. map_location: network map location as in `torch.jit.load`. more_extra_files: other extra file data names to load from bundle, see `_extra_files` of `torch.jit.load`. Returns: Triple containing loaded object, metadata dict, and extra files dict containing other file data if present cSsi|] }|dqS)rr+).0fr+r+r, wsz*load_net_with_metadata..rz{})rr$r%loaddictitemsrloads)r r.rr)r r(Zjson_data_dictr+r+r,load_net_with_metadatags  r6)TFNN)Nr+) __future__rrrr collections.abcrrtypingrrr$ monai.configr monai.utilsrrr-r6r+r+r+r, s"   M