frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Serializes a std::unique_ptr<std::mutex> object.

Browse Source 
Reviewed By: xianjiec

Differential Revision: D4901097

fbshipit-source-id: 067d6fe3e2b201818eb6967a02b0ac0289fe8327
This commit is contained in:
Bor-Yiing Su
2017-04-17 19:42:19 -07:00
committed by Facebook Github Bot
parent 0a726af42e
commit 391fd14115
2 changed files with 42 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
caffe2/sgd/iter_op.cc
View File

@ -1,10 +1,33 @@
#include "caffe2/sgd/iter_op.h"
namespace caffe2 {
void MutexSerializer::Serialize(
const Blob& blob,
const string& name,
BlobSerializerBase::SerializationAcceptor acceptor) {
CAFFE_ENFORCE(blob.IsType<std::unique_ptr<std::mutex>>());
BlobProto blob_proto;
blob_proto.set_name(name);
blob_proto.set_type("std::unique_ptr<std::mutex>");
blob_proto.set_content("");
acceptor(name, blob_proto.SerializeAsString());
}
void MutexDeserializer::Deserialize(const BlobProto& /* unused */, Blob* blob) {
*blob->GetMutable<std::unique_ptr<std::mutex>>() =
caffe2::make_unique<std::mutex>();
}
namespace {
REGISTER_CPU_OPERATOR(Iter, IterOp<CPUContext>);
REGISTER_CPU_OPERATOR(AtomicIter, AtomicIterOp<CPUContext>);
REGISTER_BLOB_SERIALIZER(
(TypeMeta::Id<std::unique_ptr<std::mutex>>()),
MutexSerializer);
REGISTER_BLOB_DESERIALIZER(std::unique_ptr<std::mutex>, MutexDeserializer);
OPERATOR_SCHEMA(Iter)
.NumInputs(0, 1)
.NumOutputs(1)

19
caffe2/sgd/iter_op.h
View File

@ -4,6 +4,7 @@
#include <limits>
#include <mutex>
#include "caffe2/core/blob_serialization.h"
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
@ -69,6 +70,24 @@ class AtomicIterOp final : public Operator<Context> {
}
};
class MutexSerializer : public BlobSerializerBase {
public:
/**
* Serializes a std::unique_ptr<std::mutex>. Note that this blob has to
* contain std::unique_ptr<std::mutex>, otherwise this function produces a
* fatal error.
*/
void Serialize(
const Blob& blob,
const string& name,
BlobSerializerBase::SerializationAcceptor acceptor) override;
};
class MutexDeserializer : public BlobDeserializerBase {
public:
void Deserialize(const BlobProto& proto, Blob* blob) override;
};
} // namespace caffe2
#endif // CAFFE2_SGD_ITER_OP_H_