mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
f73e32cd04
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/43508 Differential Revision: D22952007 fbshipit-source-id: 11e28d20175271e6068edce8cb36f9fcf867a02a
61 lines
1.6 KiB
C++
61 lines
1.6 KiB
C++
#ifndef CAFFE2_OPERATORS_LOAD_SAVE_OP_UTIL_H_
|
|
#define CAFFE2_OPERATORS_LOAD_SAVE_OP_UTIL_H_
|
|
|
|
#include <set>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
|
|
#include "caffe2/core/blob.h"
|
|
#include "caffe2/core/blob_serialization.h"
|
|
|
|
namespace caffe2 {
|
|
namespace load_save_op_util {
|
|
|
|
struct BlobState {
|
|
int64_t total_size;
|
|
int64_t current_size;
|
|
bool is_tensor;
|
|
std::set<int32_t> seen_chunks_ids;
|
|
|
|
explicit BlobState(
|
|
int64_t total_size = 0,
|
|
int64_t current_size = 0,
|
|
bool is_tensor = false)
|
|
: total_size(total_size),
|
|
current_size(current_size),
|
|
is_tensor(is_tensor) {}
|
|
};
|
|
|
|
CAFFE2_API std::string buildBlobNameFromDbKey(
|
|
const std::string& dbKey,
|
|
const std::string& strip_prefix = "",
|
|
const std::string& add_prefix = "");
|
|
|
|
// We are tracking sizes of already read tensor parts while reading data
|
|
// chunks. This way we can make sure that all chunks were loaded in the end.
|
|
CAFFE2_API void ProcessBlob(
|
|
Blob* blob,
|
|
const BlobProto& proto,
|
|
std::unordered_map<std::string, BlobState>* blob_states_ptr,
|
|
const std::string& key,
|
|
int* loaded_blobs);
|
|
|
|
CAFFE2_API void prepareBlob(
|
|
Blob* blob,
|
|
std::unordered_map<std::string, BlobState>* blob_states_ptr,
|
|
const std::string& key);
|
|
|
|
CAFFE2_API void updateBlobStates(
|
|
const BlobProto& proto,
|
|
std::unordered_map<std::string, BlobState>* blob_states_ptr,
|
|
const std::string& key,
|
|
int* loaded_blobs);
|
|
|
|
CAFFE2_API void validateBlobStates(
|
|
const std::unordered_map<std::string, BlobState>& blob_states);
|
|
|
|
} // namespace load_save_op_util
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_OPERATORS_LOAD_SAVE_OP_UTIL_H_
|