Add train() / eval() / is_training() to C++ ScriptModule API #16044
Conversation
facebook-github-bot
added
the
oncall: jit
yf225
changed the title
yf225
requested review from
gchanan,
goldsborough and
zdevito
and removed request for
goldsborough
test/cpp/jit/tests.h
Outdated
| void testEvalModeForLoadedModule() { | ||
| std::string module_path = "dropout_model.pt"; | ||
| std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(module_path); | ||
| // Test eval mode |
There was a problem hiding this comment.
I think the comments here don't add too much
torch/csrc/jit/script/module.h
Outdated
| } | ||
| void train(bool on = true) { | ||
| for (auto& submod : get_modules()) { | ||
| submod.value().module->train(on); |
There was a problem hiding this comment.
you can also write submod->module->train(on)
torch/csrc/jit/script/module.h
Outdated
| for (auto& submod : get_modules()) { | ||
| submod.value().module->train(on); | ||
| } | ||
| auto t = autograd::make_variable(at::full({}, on ? 1 : 0, at::kLong)); |
There was a problem hiding this comment.
You can replace
auto t = autograd::make_variable(at::full({}, on ? 1 : 0, at::kLong));
with
auto t = torch::tensor(on ? 1 : 0, at::kLong);
To clarify: torch:: factory functions create variables while at:: functions create tensors. torch::tensor/at::tensor is like torch.tensor in python (creates a tensor with the values you give it)`
There was a problem hiding this comment.
And then just embed it in the register_parameter call:
register_parameter("training", torch::tensor(on ? 1 : 0, at::kLong), /*is_buffer=*/true);
torch/csrc/jit/script/module.h
Outdated
| } | ||
| bool is_training() { | ||
| if (auto p = find_parameter("training")) { | ||
| return (*p->slot()).item().toLong() == 1; |
There was a problem hiding this comment.
Here you are converting the tensor first to a Scalar, and then to a long. You can just use item<T> to get the value of the tensor without the intermediate Scalar:
(*p->slot()).item<int64_t>()
torch/csrc/jit/script/module.h
Outdated
torch/csrc/jit/script/module.h
Outdated
There was a problem hiding this comment.
nit: remove else block, just return true + comment (since you return in the if block)
There was a problem hiding this comment.
@yf225 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
There was a problem hiding this comment.
@yf225 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
There was a problem hiding this comment.
@yf225 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
yf225
force-pushed
the
scriptmodule_eval
branch
from
82be773 to
4c36017
Compare
yf225
force-pushed
the
scriptmodule_eval
branch
from
f5a442d to
c03b1e5
Compare
There was a problem hiding this comment.
@yf225 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
4 participants
This PR aims to fix https://discuss.pytorch.org/t/how-to-change-a-loaded-model-to-evaluation-mode-in-c/32330, by adding
train()/eval()/is_training()to C++ ScriptModule API.