pytorch/torch/csrc/jit/api/object.cpp
Jerry Zhang 6fa76b8a0c [jit] __deepcopy__ for RecursiveScriptModule (#32684)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/32684

Previously we have `clone` and `clone_instance`, where `clone` will clone both type
and value, and `clone_instance` only clone the value, both of them are shallow copies.
We need to re-evaluate whether we should expose them as a user facing API.
I think we should hide `clone`, but `clone_instance` might be useful as well, especially
when we are copying a model with very large weights, people might just want to do shallow copy.

This PR adds a `deepcopy` that might be useful as a user API, which deep copies the values, including
Tensor, but we didn't deepcopy `Blob`, `Capsule`, `Future` or `PyObject`.
For more discussions please see the following issue.

fixes: https://github.com/pytorch/pytorch/issues/32519

Test Plan: Imported from OSS

Differential Revision: D21220756

fbshipit-source-id: 476bf11fe82c08fac36e7457879a09f545ffdc5e
2020-04-28 18:47:11 -07:00

58 lines
1.5 KiB
C++

#include <torch/csrc/jit/api/object.h>
#include <ATen/core/jit_type.h>
#include <torch/csrc/jit/api/compilation_unit.h>
#include <torch/csrc/jit/frontend/resolver.h>
#include <torch/csrc/jit/frontend/sugared_value.h>
namespace torch {
namespace jit {
Object::Object(
std::shared_ptr<CompilationUnit> cu,
const c10::ClassTypePtr& type)
: Object(c10::ivalue::Object::create(
c10::StrongTypePtr(std::move(cu), type),
type->numAttributes())) {}
ObjectPtr Object::_ivalue() const {
TORCH_INTERNAL_ASSERT(_ivalue_);
return _ivalue_;
}
c10::optional<Method> Object::find_method(const std::string& basename) const {
for (Function* fn : type()->methods()) {
if (fn->name() == basename) {
return Method(_ivalue(), fn);
}
}
return c10::nullopt;
}
void Object::define(const std::string& src, const ResolverPtr& resolver) {
const auto self = SimpleSelf(type());
_ivalue()->compilation_unit()->define(
*type()->name(), src, resolver ? resolver : nativeResolver(), &self);
}
Object Object::deepcopy() const {
c10::IValue::HashAliasedIValueMap memo;
return deepcopy(memo);
}
Object Object::deepcopy(c10::IValue::HashAliasedIValueMap& memo) const {
Object obj(_ivalue()->compilation_unit(), type());
// Deepcopy slots. If a slot is a module - recursively copy it.
size_t N = type()->numAttributes();
for (size_t i = 0; i < N; ++i) {
IValue s = _ivalue()->getSlot(i);
obj._ivalue()->setAttr(type()->getAttributeName(i), s.deepcopy(memo));
}
return obj;
}
} // namespace jit
} // namespace torch