Google Git
Sign in
android / platform / external / pytorch / 598b713660 / . / torch / csrc / jit / autodiff.cpp
blob: fdfa031a0e240475d41e807a31b15a15389f5a89 [file] [log] [blame]
#include <torch/csrc/jit/autodiff.h>
#include <torch/csrc/jit/operator.h>
#include <torch/csrc/jit/passes/common_subexpression_elimination.h>
#include <torch/csrc/jit/passes/constant_pooling.h>
#include <torch/csrc/jit/passes/dead_code_elimination.h>
#include <torch/csrc/jit/passes/lower_tuples.h>
#include <torch/csrc/jit/script/compiler.h>
#include <torch/csrc/jit/symbolic_script.h>
#include <torch/csrc/jit/symbolic_variable.h>
#include <torch/csrc/utils/functional.h>
#include <c10/util/Exception.h>
#include <algorithm>
#include <memory>
namespace torch {
namespace jit {
using value_map = std::unordered_map<Value*, Value*>;
using value_set = std::unordered_set<Value*>;
void wrapDim(int64_t& dim, const std::vector<int64_t>& sizes) {
if (dim < 0) {
dim += sizes.size();
}
}
bool isDifferentiable(Node* n) {
// TODO: scalar-tensor ops should be canonicalized
static OperatorSet differentiable_ops = {
"aten::add(Tensor self, Tensor other, *, Scalar alpha) -> Tensor",
"aten::add(Tensor self, Scalar other, Scalar alpha) -> Tensor",
"aten::sub(Tensor self, Tensor other, *, Scalar alpha) -> Tensor",
"aten::sub(Tensor self, Scalar other, Scalar alpha) -> Tensor",
"aten::mul(Tensor self, Tensor other) -> Tensor",
"aten::mul(Tensor self, Scalar other) -> Tensor",
"aten::div(Tensor self, Tensor other) -> Tensor",
"aten::div(Tensor self, Scalar other) -> Tensor",
"aten::max(Tensor self, Tensor other) -> Tensor",
"aten::min(Tensor self, Tensor other) -> Tensor",
"aten::sigmoid(Tensor self) -> Tensor",
"aten::tanh(Tensor self) -> Tensor",
"aten::relu(Tensor self) -> Tensor",
"aten::threshold(Tensor self, Scalar threshold, Scalar value) -> Tensor",
"aten::erf(Tensor self) -> Tensor",
"aten::erfc(Tensor self) -> Tensor",
"aten::exp(Tensor self) -> Tensor",
"aten::t(Tensor self) -> Tensor",
"aten::neg(Tensor self) -> Tensor",
"aten::clamp(Tensor self, Scalar? min, Scalar? max) -> Tensor",
"aten::where(Tensor condition, Tensor self, Tensor other) -> Tensor",
"aten::type_as(Tensor self, Tensor other) -> Tensor",
"aten::unsqueeze(Tensor self, int dim) -> Tensor",
"aten::addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta, Scalar alpha) -> Tensor",
"aten::mm(Tensor self, Tensor mat2) -> Tensor",
"aten::lt(Tensor self, Tensor other) -> Tensor",
"aten::le(Tensor self, Tensor other) -> Tensor",
"aten::gt(Tensor self, Tensor other) -> Tensor",
"aten::ge(Tensor self, Tensor other) -> Tensor",
"aten::eq(Tensor self, Tensor other) -> Tensor",
"aten::ne(Tensor self, Tensor other) -> Tensor",
"aten::lt(Tensor self, Scalar other) -> Tensor",
"aten::le(Tensor self, Scalar other) -> Tensor",
"aten::gt(Tensor self, Scalar other) -> Tensor",
"aten::ge(Tensor self, Scalar other) -> Tensor",
"aten::eq(Tensor self, Scalar other) -> Tensor",
"aten::ne(Tensor self, Scalar other) -> Tensor",
"aten::abs(Tensor self) -> Tensor",
"aten::acos(Tensor self) -> Tensor",
"aten::asin(Tensor self) -> Tensor",
"aten::atan(Tensor self) -> Tensor",
"aten::ceil(Tensor self) -> Tensor",
"aten::cos(Tensor self) -> Tensor",
"aten::cosh(Tensor self) -> Tensor",
"aten::exp(Tensor self) -> Tensor",
"aten::expm1(Tensor self) -> Tensor",
"aten::floor(Tensor self) -> Tensor",
"aten::fmod(Tensor self, Scalar other) -> Tensor",
"aten::frac(Tensor self) -> Tensor",
"aten::log(Tensor self) -> Tensor",
"aten::log10(Tensor self) -> Tensor",
"aten::log1p(Tensor self) -> Tensor",
"aten::log2(Tensor self) -> Tensor",
"aten::reciprocal(Tensor self) -> Tensor",
"aten::remainder(Tensor self, Scalar other) -> Tensor",
"aten::round(Tensor self) -> Tensor",
"aten::rsqrt(Tensor self) -> Tensor",
"aten::sin(Tensor self) -> Tensor",
"aten::sinh(Tensor self) -> Tensor",
"aten::tan(Tensor self) -> Tensor",
"aten::trunc(Tensor self) -> Tensor",
"prim::SumToSize(Tensor(a) self, int[] size) -> Tensor(a)",
"aten::log_softmax(Tensor self, int dim) -> Tensor",
"aten::avg_pool2d(Tensor self, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, bool count_include_pad) -> Tensor",
"aten::max_pool2d_with_indices(Tensor self, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode) -> (Tensor, Tensor)",
"aten::thnn_conv2d_forward(Tensor self, Tensor weight, int[] kernel_size, Tensor? bias, int[] stride, int[] padding) -> (Tensor, Tensor, Tensor)",
"aten::native_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor)",
};
// TODO: add support for the following fusible operators.
// They're a little tricky to implement; max/min require mutability for best
// perf "aten::atan2(Tensor self) -> Tensor", "aten::max(Tensor self) ->
// Tensor", "aten::min(Tensor self) -> Tensor"
if (n->kind() == prim::Constant || n->kind() == prim::Undefined ||
n->kind() == prim::AutogradAdd || n->kind() == prim::ConstantChunk ||
n->kind() == prim::None)
return true;
if (differentiable_ops.find(n))
return true;
auto schema = n->maybeSchema();
if (schema && hasGradientInfoForSchema(*schema)) {
return true;
}
if (n->matches(
"aten::expand(Tensor self, int[] size, *, bool implicit) -> Tensor")) {
return n->get<std::vector<int64_t>>(attr::size) &&
n->is_constant(attr::implicit) &&
n->namedInput(attr::self)->type()->cast<CompleteTensorType>();
}
if (n->matches("aten::view(Tensor self, int[] size) -> Tensor")) {
return n->get<std::vector<int64_t>>(attr::size) &&
n->namedInput(attr::self)->type()->cast<CompleteTensorType>();
}
if (n->matches(
"aten::nll_loss(Tensor self, Tensor target, Tensor? weight, int reduction, int ignore_index) -> Tensor")) {
// TODO(asuhan): support weight
return n->namedInput(attr::weight)->node()->kind() == prim::Undefined;
}
// linear blocks may appear as inputs to graph executors, but they are removed
// before differentiation occurs
if (n->kind() == prim::GradOf) {
auto body = n->blocks().at(0);
return std::all_of(
body->nodes().begin(),
body->nodes().end(),
static_cast<bool (*)(Node*)>(isDifferentiable));
}
return false;
}
bool isDifferentiable(Graph& g) {
return std::all_of(
g.nodes().begin(),
g.nodes().end(),
static_cast<bool (*)(Node*)>(isDifferentiable));
}
// NB: Write gradient using torchscript
// For example, node aten::mul() should be defined as follows
// def forward(x, y):
// return x*y, (x, y)
// def backward(ctx, grad_output):
// x, y = ctx
// return (y * grad_output).sum_to_size(x), (x * grad_output).sum_to_size(y)
//
// Here ctx is a tuple that carries all input/intermediate results needed in
// backward from forward pass.
//
// This python code is compiled into a GradientPair which includes a forward
// graph and a backward graph. Forward graph will be used to replace the node in
// grad_desc.f, and backward graph will be used to construct GradOf(node) in
// reverse_block. Grad_values(a.k.a gradOutputs) propagated through
// node->owningGraph() in **reversed** order, thus GradientPair.forward ahould
// be inserted **after** the node being replaced, so that we don't traverse the
// graph infinite times.
//
// The output of compiled forward graph is [real_outputs, ctx]
// The input of compiled backward graph is [ctx, grad_values]
// We run LowerSimpleTuples afterwards to elmininate all tuples generated in
// this process. The original node and TupleConstruct nodes in forward graph
// will be cleaned up later using EliminateDeadCode(block). TupleUnPack node in
// backward graph will be removed in eliminateDeadcode(ReverseDetails) defined
// in this file.
static c10::optional<std::vector<Value*>> build_script_grad(
Node* node,
const ArrayRef<Value*>& grads) {
auto graph = node->owningGraph();
auto compiled_graphs = gradientInfoForSchema(node->schema());
if (!compiled_graphs) {
return c10::nullopt;
}
// Use forward graph to replace node in grad_desc.f
value_list new_outputs;
{
WithInsertPoint guard(node->next());
auto fw_graph = compiled_graphs->forward;
new_outputs = inlineCallTo(
*graph, *fw_graph, node->inputs(), /*unpack_outputs=*/true);
for (size_t i = 0; i < node->outputs().size(); ++i) {
new_outputs.at(i)->setType(node->outputs()[i]->type());
new_outputs.at(i)->replaceAllUsesWith(node->outputs()[i]);
}
}
// Use backward graph to construct reverse_block
auto bw_graph = compiled_graphs->backward;
auto grad_vec = grads.vec();
auto it = grad_vec.begin();
grad_vec.insert(it, new_outputs.back());
ArrayRef<Value*> grad(grad_vec);
auto grad_inputs =
inlineCallTo(*graph, *bw_graph, grad, /*unpack_outputs=*/true);
return grad_inputs;
};
namespace {
class GradientHelper {
public:
GradientHelper(Node* n) : node(n) {}
std::vector<Value*> gradient(ArrayRef<Value*> grad_values) {
if (!isDifferentiable(node)) {
throw std::runtime_error(
std::string("differentiation of ") + node->kind().toDisplayString() +
" is not supported, or it is missing necessary type information");
}
// If AD is defined using torchscript, use it instead of symbolic
auto script_grads = build_script_grad(node, grad_values);
if (script_grads)
return *script_grads;
// Definition not found in torchscript, look up in the buildSymbolicGradient
// TODO: migrate all to using torchscript
auto sym_grads = buildSymbolicGradient(fmap<SymbolicVariable>(grad_values));
return fmap(sym_grads, [](const SymbolicVariable& v) { return v.value(); });
}
private:
Node* node;
SymbolicVariable sumToSizeOf(SymbolicVariable v, Symbol input_name) {
Value* size;
{
WithInsertPoint insert_guard{node};
size = SymbolicVariable(node->namedInput(input_name)).size();
}
return v.sumToSize(size);
};
std::vector<SymbolicVariable> buildSymbolicGradient(
const std::vector<SymbolicVariable>& grads) {
static const OperatorSet comparison_ops = {
"aten::lt(Tensor self, Tensor other) -> Tensor",
"aten::le(Tensor self, Tensor other) -> Tensor",
"aten::gt(Tensor self, Tensor other) -> Tensor",
"aten::ge(Tensor self, Tensor other) -> Tensor",
"aten::eq(Tensor self, Tensor other) -> Tensor",
"aten::ne(Tensor self, Tensor other) -> Tensor",
"aten::lt(Tensor self, Scalar other) -> Tensor",
"aten::le(Tensor self, Scalar other) -> Tensor",
"aten::gt(Tensor self, Scalar other) -> Tensor",
"aten::ge(Tensor self, Scalar other) -> Tensor",
"aten::eq(Tensor self, Scalar other) -> Tensor",
"aten::ne(Tensor self, Scalar other) -> Tensor",
};
auto inputs = fmap<SymbolicVariable>(node->inputs());
auto outputs = fmap<SymbolicVariable>(node->outputs());
if (node->matches(
"aten::add(Tensor self, Tensor other, *, Scalar alpha) -> Tensor")) {
return {
sumToSizeOf(grads.at(0), attr::self),
sumToSizeOf(grads.at(0) * node->namedInput(attr::alpha), attr::other),
nullptr};
} else if (
node->matches(
"aten::add(Tensor self, Scalar other, Scalar alpha) -> Tensor")) {
return {grads.at(0), nullptr, nullptr};
} else if (node->kind() == prim::AutogradAdd) {
// NB: AutogradAdds don't broadcast
return {grads.at(0), grads.at(0)};
} else if (
node->matches(
"aten::sub(Tensor self, Tensor other, *, Scalar alpha) -> Tensor")) {
return {sumToSizeOf(grads.at(0), attr::self),
sumToSizeOf(
-grads.at(0) * node->namedInput(attr::alpha), attr::other),
nullptr};
} else if (
node->matches(
"aten::sub(Tensor self, Scalar other, Scalar alpha) -> Tensor")) {
return {grads.at(0), nullptr, nullptr};
} else if (node->matches(
"aten::mul(Tensor self, Tensor other) -> Tensor")) {
return {sumToSizeOf(grads.at(0) * inputs.at(1), attr::self),
sumToSizeOf(grads.at(0) * inputs.at(0), attr::other)};
} else if (node->matches(
"aten::mul(Tensor self, Scalar other) -> Tensor")) {
return {grads.at(0) * inputs.at(1), nullptr};
} else if (node->matches(
"aten::div(Tensor self, Tensor other) -> Tensor")) {
return {sumToSizeOf(grads.at(0) / inputs.at(1), attr::self),
sumToSizeOf(
-grads.at(0) * inputs.at(0) / (inputs.at(1) * inputs.at(1)),
attr::other)};
} else if (node->matches(
"aten::div(Tensor self, Scalar other) -> Tensor")) {
return {grads.at(0) / inputs.at(1), nullptr};
} else if (node->matches(
"aten::max(Tensor self, Tensor other) -> Tensor")) {
return {
sumToSizeOf(
grads.at(0) * (inputs.at(0) > inputs.at(1)).type_as(grads.at(0)),
attr::self),
sumToSizeOf(
grads.at(0) * (inputs.at(1) > inputs.at(0)).type_as(grads.at(0)),
attr::other)};
} else if (node->matches(
"aten::min(Tensor self, Tensor other) -> Tensor")) {
return {
sumToSizeOf(
grads.at(0) * (inputs.at(0) < inputs.at(1)).type_as(grads.at(0)),
attr::self),
sumToSizeOf(
grads.at(0) * (inputs.at(1) < inputs.at(0)).type_as(grads.at(0)),
attr::other)};
} else if (
node->matches(
"aten::where(Tensor condition, Tensor self, Tensor other) -> Tensor")) {
return {nullptr,
sumToSizeOf(
grads.at(0) * inputs.at(0).type_as(grads.at(0)), attr::self),
sumToSizeOf(
grads.at(0) * (1 - inputs.at(0)).type_as(grads.at(0)),
attr::other)};
} else if (node->matches("aten::sigmoid(Tensor self) -> Tensor")) {
// TODO: The order of operations matter in this case. This
// works for ppc64le and x86_64. Need to look at why the
// order matters.
return {(1 - outputs.at(0)) * outputs.at(0) * grads.at(0)};
} else if (node->matches("aten::tanh(Tensor self) -> Tensor")) {
return {grads.at(0) * (1 - outputs.at(0) * outputs.at(0))};
} else if (node->matches("aten::relu(Tensor self) -> Tensor")) {
return {grads.at(0) *
(outputs.at(0) > at::Scalar(0)).type_as(outputs.at(0))};
} else if (
node->matches(
"aten::clamp(Tensor self, Scalar? min, Scalar? max) -> Tensor")) {
// handle the case that min/max is None
Value* min = inputs.at(1);
bool min_must_be_none = min->node()->kind() == prim::None;
Value* max = inputs.at(2);
bool max_must_be_none = max->node()->kind() == prim::None;
// XXX - this formula is wrong when min or max are not stricly prim::None
// but may be None dynamically. In this case an internal compiler error
// will get thrown when trying to generate expressions involving the
// values of min/max
if (!min_must_be_none && !max_must_be_none) {
return {grads.at(0) *
(1 - (inputs.at(0) <= inputs.at(1)).type_as(inputs.at(0))) *
(1 - (inputs.at(0) >= inputs.at(2)).type_as(inputs.at(0))),
nullptr,
nullptr};
} else if (max_must_be_none) {
return {grads.at(0) *
(1 - (inputs.at(0) <= inputs.at(1)).type_as(inputs.at(0))),
nullptr,
nullptr};
} else if (min_must_be_none) {
return {grads.at(0) *
(1 - (inputs.at(0) >= inputs.at(2)).type_as(inputs.at(0))),
nullptr,
nullptr};
} else {
return {grads.at(0), nullptr, nullptr};
}
} else if (
node->matches(
"aten::threshold(Tensor self, Scalar threshold, Scalar value) -> Tensor")) {
auto threshold = node->get<at::Scalar>(attr::threshold).value();
return {grads.at(0) * (inputs.at(0) > threshold).type_as(outputs.at(0)),
nullptr,
nullptr};
} else if (node->matches("aten::erf(Tensor self) -> Tensor")) {
return {grads.at(0) * 1.12837916709551 *
(-inputs.at(0) * inputs.at(0)).exp()};
} else if (node->matches("aten::erfc(Tensor self) -> Tensor")) {
return {-grads.at(0) * 1.12837916709551 *
(-inputs.at(0) * inputs.at(0)).exp()};
} else if (node->matches("aten::exp(Tensor self) -> Tensor")) {
return {grads.at(0) * (outputs.at(0))};
} else if (node->matches("aten::t(Tensor self) -> Tensor")) {
return {grads.at(0).t()};
} else if (node->matches("aten::neg(Tensor self) -> Tensor")) {
return {-grads.at(0)};
} else if (node->matches("aten::abs(Tensor self) -> Tensor")) {
return {grads.at(0) * inputs.at(0).sign()};
} else if (node->matches("aten::acos(Tensor self) -> Tensor")) {
return {grads.at(0) *
-((-inputs.at(0) * inputs.at(0) + at::Scalar(1)).rsqrt())};
} else if (node->matches("aten::asin(Tensor self) -> Tensor")) {
return {grads.at(0) *
(-inputs.at(0) * inputs.at(0) + at::Scalar(1)).rsqrt()};
} else if (node->matches("aten::atan(Tensor self) -> Tensor")) {
return {grads.at(0) / (inputs.at(0) * inputs.at(0) + at::Scalar(1))};
} else if (
node->matches(
"prim::SumToSize(Tensor(a) self, int[] size) -> Tensor(a)")) {
Value* self_size;
{
WithInsertPoint insert_guard{node};
self_size = inputs.at(0).size();
}
return {grads.at(0).expand(self_size), nullptr};
} else if (node->matches("aten::ceil(Tensor self) -> Tensor")) {
return {SymbolicVariable::zeros_like(grads.at(0))};
} else if (node->matches("aten::cos(Tensor self) -> Tensor")) {
return {grads.at(0) * -inputs.at(0).sin()};
} else if (node->matches("aten::cosh(Tensor self) -> Tensor")) {
return {grads.at(0) * inputs.at(0).sinh()};
} else if (node->matches("aten::exp(Tensor self) -> Tensor")) {
return {grads.at(0) * outputs.at(0)};
} else if (node->matches("aten::expm1(Tensor self) -> Tensor")) {
return {grads.at(0) * (outputs.at(0) + at::Scalar(1))};
} else if (node->matches("aten::floor(Tensor self) -> Tensor")) {
return {SymbolicVariable::zeros_like(grads.at(0))};
} else if (node->matches(
"aten::fmod(Tensor self, Scalar other) -> Tensor")) {
return {grads.at(0), nullptr};
} else if (node->matches("aten::frac(Tensor self) -> Tensor")) {
return {grads.at(0)};
} else if (node->matches("aten::log(Tensor self) -> Tensor")) {
return {grads.at(0) / inputs.at(0)};
} else if (node->matches("aten::log10(Tensor self) -> Tensor")) {
return {grads.at(0) / (inputs.at(0) * 2.3025850929940456)};
} else if (node->matches("aten::log1p(Tensor self) -> Tensor")) {
return {grads.at(0) / (inputs.at(0) + at::Scalar(1))};
} else if (node->matches("aten::log2(Tensor self) -> Tensor")) {
return {grads.at(0) / (inputs.at(0) * 0.6931471805599453)};
} else if (node->matches("aten::reciprocal(Tensor self) -> Tensor")) {
return {-grads.at(0) * outputs.at(0) * outputs.at(0)};
} else if (node->matches(
"aten::remainder(Tensor self, Scalar other) -> Tensor")) {
return {grads.at(0), nullptr};
} else if (node->matches("aten::round(Tensor self) -> Tensor")) {
return {SymbolicVariable::zeros_like(grads.at(0))};
} else if (node->matches("aten::rsqrt(Tensor self) -> Tensor")) {
return {grads.at(0) * outputs.at(0).pow(3.) * -0.5};
} else if (node->matches("aten::sin(Tensor self) -> Tensor")) {
return {grads.at(0) * inputs.at(0).cos()};
} else if (node->matches("aten::sinh(Tensor self) -> Tensor")) {
return {grads.at(0) * inputs.at(0).cosh()};
} else if (node->matches("aten::tan(Tensor self) -> Tensor")) {
return {grads.at(0) * (1. + outputs.at(0) * outputs.at(0))};
} else if (node->matches("aten::trunc(Tensor self) -> Tensor")) {
return {SymbolicVariable::zeros_like(grads.at(0))};
} else if (node->kind() == prim::ConstantChunk) {
return {SymbolicVariable::cat(grads, node->i(attr::dim))};
} else if (
node->matches("aten::view(Tensor self, int[] size) -> Tensor") ||
node->matches("aten::reshape(Tensor self, int[] shape) -> Tensor")) {
// TODO: if sizes are not available statically, add an operator that
// reutrns them as a tuple
auto sizes = node->namedInput(attr::self)
->type()
->expect<CompleteTensorType>()
->sizes();
return {grads.at(0).reshape(sizes), nullptr};
} else if (node->matches(
"aten::type_as(Tensor self, Tensor other) -> Tensor")) {
return {grads.at(0).type_as(inputs.at(0)), nullptr};
} else if (node->matches(
"aten::unsqueeze(Tensor self, int dim) -> Tensor")) {
return {grads.at(0).squeeze(node->namedInput(attr::dim)), nullptr};
} else if (
node->matches(
"aten::addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta, Scalar alpha) -> Tensor")) {
return {
sumToSizeOf(grads.at(0) * node->namedInput(attr::beta), attr::self),
grads.at(0).mm(inputs.at(2).t()) * node->namedInput(attr::alpha),
inputs.at(1).t().mm(grads.at(0)) * node->namedInput(attr::alpha),
nullptr,
nullptr};
} else if (node->matches("aten::mm(Tensor self, Tensor mat2) -> Tensor")) {
return {grads.at(0).mm(inputs.at(1).t()),
inputs.at(0).t().mm(grads.at(0))};
} else if (
node->matches(
"aten::expand(Tensor self, int[] size, *, bool implicit) -> Tensor")) {
const auto& input_sizes = inputs.at(0).sizes();
if (input_sizes.size() == 0)
return {grads.at(0).sum(), nullptr, nullptr};
auto grad_sizes = node->get<std::vector<int64_t>>(attr::size).value();
auto grad = grads.at(0);
while (grad_sizes.size() > input_sizes.size()) {
grad = grad.sum(0, false);
grad_sizes.erase(grad_sizes.begin());
}
for (size_t i = 0; i < input_sizes.size(); ++i) {
if (input_sizes[i] == 1 && grad_sizes[i] > 1) {
grad = grad.sum(i, true);
}
}
return {grad, nullptr, nullptr};
} else if (node->matches("aten::squeeze(Tensor self) -> Tensor")) {
const auto& sizes = inputs.at(0).sizes();
std::vector<size_t> squeezed_dims;
for (size_t i = 0; i < sizes.size(); ++i) {
if (sizes[i] != 1)
continue;
squeezed_dims.push_back(i);
}
SymbolicVariable returned_grad = grads.at(0);
for (const auto& dim : squeezed_dims) {
returned_grad = returned_grad.unsqueeze(dim);
}
return {returned_grad};
} else if (node->matches(
"aten::squeeze(Tensor self, int dim) -> Tensor",
/*const_inputs=*/attr::dim)) {
int64_t dim = *node->get<int64_t>(attr::dim);
const auto& sizes = inputs.at(0).sizes();
wrapDim(dim, sizes);
if (sizes.size() == 0) {
return {grads.at(0), nullptr};
}
return {sizes.at(dim) > 1 ? grads.at(0) : grads.at(0).unsqueeze(dim),
nullptr};
} else if (node->matches(
"aten::cat(Tensor[] tensors, int dim) -> Tensor",
/*const_inputs=*/attr::dim)) {
int dim = *node->get<int64_t>(attr::dim);
auto tensor_inputs = inputs;
tensor_inputs.pop_back();
const auto& first_sizes = tensor_inputs.at(0).sizes();
const auto has_first_sizes = [&first_sizes](SymbolicVariable var) {
return var.sizes() == first_sizes;
};
// NB: this is a specialization for the common case where all inputs are
// of equal sizes. We can use a single split operation to handle that.
if (std::all_of(
tensor_inputs.begin(), tensor_inputs.end(), has_first_sizes)) {
auto tensor_grads = grads.at(0).chunk(tensor_inputs.size(), dim);
tensor_grads.emplace_back(nullptr); // for attr::dim
return tensor_grads;
} else {
size_t offset = 0;
auto grad = grads.at(0);
std::vector<SymbolicVariable> tensor_grads;
for (auto input : tensor_inputs) {
tensor_grads.push_back(grad.narrow(dim, offset, input.sizes()[dim]));
offset += input.sizes()[dim];
}
tensor_grads.emplace_back(nullptr); // for attr::dim
return tensor_grads;
}
} else if (comparison_ops.find(node)) {
return {nullptr, nullptr};
} else if (
node->matches(
"aten::avg_pool2d(Tensor self, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, bool count_include_pad) -> Tensor")) {
AT_ASSERT(grads.size() == 1);
auto graph = node->owningGraph();
auto backward_value = graph->insert(
aten::avg_pool2d_backward,
{grads.at(0).value(),
node->namedInput(attr::self),
node->namedInput(attr::kernel_size),
node->namedInput(attr::stride),
node->namedInput(attr::padding),
node->namedInput(attr::ceil_mode),
node->namedInput(attr::count_include_pad)});
return {backward_value->node()->output(0),
nullptr,
nullptr,
nullptr,
nullptr,
nullptr};
} else if (
node->matches(
"aten::max_pool2d_with_indices(Tensor self, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode) -> (Tensor, Tensor)")) {
AT_ASSERT(grads.size() == 2);
auto graph = node->owningGraph();
auto backward_value = graph->insert(
aten::max_pool2d_with_indices_backward,
{grads.at(0).value(),
node->namedInput(attr::self),
node->namedInput(attr::kernel_size),
node->namedInput(attr::stride),
node->namedInput(attr::padding),
node->namedInput(attr::dilation),
node->namedInput(attr::ceil_mode),
outputs.at(1).value()});
return {backward_value->node()->output(0),
nullptr,
nullptr,
nullptr,
nullptr,
nullptr};
} else if (
node->matches(
"aten::thnn_conv2d_forward(Tensor self, Tensor weight, int[] kernel_size, Tensor? bias, int[] stride, int[] padding) -> (Tensor, Tensor, Tensor)")) {
auto graph = node->owningGraph();
auto backward_value = graph->insert(
aten::thnn_conv2d_backward,
{grads.at(0).value(),
inputs.at(0).value(),
inputs.at(1).value(),
node->namedInput(attr::kernel_size),
node->namedInput(attr::stride),
node->namedInput(attr::padding),
outputs.at(1).value(),
outputs.at(2).value(),
graph->insertConstant(std::vector<bool>{true, true, true})});
// graph->insert returns a tuple automatically if multiple outputs are
// returned. So unpack them again.
Node* tuple_unpack_node =
graph->insertNode(graph->createTupleUnpack(backward_value));
auto tuple_outputs = tuple_unpack_node->outputs();
AT_ASSERT(tuple_outputs.size() == size_t(3));
return {tuple_outputs[0],
tuple_outputs[1],
nullptr,
tuple_outputs[2],
nullptr,
nullptr};
} else if (
node->matches(
"aten::native_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor)")) {
auto graph = node->owningGraph();
auto backward_value = graph->insert(
aten::native_batch_norm_backward,
{grads.at(0).value(),
inputs.at(0).value(),
inputs.at(1).value(),
inputs.at(3).value(),
inputs.at(4).value(),
outputs.at(1).value(),
outputs.at(2).value(),
inputs.at(5).value(),
inputs.at(7).value(),
graph->insertConstant(std::vector<bool>{true, true, true})});
// graph->insert returns a tuple automatically if multiple outputs are
// returned. So unpack them again.
Node* tuple_unpack_node =
graph->insertNode(graph->createTupleUnpack(backward_value));
auto tuple_outputs = tuple_unpack_node->outputs();
AT_ASSERT(tuple_outputs.size() == size_t(3));
return {tuple_outputs[0],
tuple_outputs[1],
tuple_outputs[2],
nullptr,
nullptr,
nullptr,
nullptr,
nullptr};
} else if (
node->matches(
"aten::nll_loss(Tensor self, Tensor target, Tensor? weight, int reduction, int ignore_index) -> Tensor")) {
auto graph = node->owningGraph();
auto total_weight = graph->insertNode(graph->createUndefined());
auto weight = graph->insertNode(graph->createUndefined());
auto backward_value = graph->insert(
aten::nll_loss_backward,
{grads.at(0).value(),
inputs.at(0).value(),
inputs.at(1).value(),
weight->output(),
inputs.at(3).value(),
inputs.at(4).value(),
total_weight->output()});
return {backward_value->node()->output(0),
nullptr,
nullptr,
nullptr,
nullptr};
} else if (node->matches(
"aten::log_softmax(Tensor self, int dim) -> Tensor")) {
AT_ASSERT(grads.size() == 1);
auto graph = node->owningGraph();
auto backward_value = graph->insert(
aten::_log_softmax_backward_data,
{grads.at(0).value(),
outputs.at(0).value(),
node->namedInput(attr::dim),
node->namedInput(attr::self)});
return {backward_value->node()->output(0), nullptr};
} else if (
node->kind() == prim::Constant || node->kind() == prim::Undefined ||
node->kind() == prim::None) {
return {};
}
throw std::runtime_error(
std::string("failed to differentiate `") +
node->kind().toDisplayString() + "`");
}
};
} // namespace
// If we have a function y = f(x) with jacobian J, the backwards of f is dx =
// J^t dy. Note that because the backwards always implements this matrix
// multiply, we know that it maps an input vector of zeros to an output vector
// of zero regardless of what operations it choses to do inside to actually
// implement the matrix multiply (most use some optimized form and never
// generate J^t). More generally, we know that all of the backward computations
// are linear and can use this property to do more aggressive optimizations
// later. It is ok to replace any backward function with known-zero inputs with
// something that produces known-zero outputs. This function encloses each
// know-linear backward function in a 'GradOf' sub-block so that we can perform
// optimizations using this information. In particular, specializeUndef will
// observe if all the inputs to the linear block are Undef, which the autograd
// uses to represent zeros, and then propagate the undefs to the outputs of the
// block.
static std::vector<Value*> linearGradientForNode(
Node* node,
ArrayRef<Value*> grad_values) {
auto& graph = *node->owningGraph();
auto linear = graph.insertNode(graph.create(prim::GradOf, {grad_values}, 0));
// to make reading gradient graphs easier, remember the name of the forward op
linear->s_(attr::name, node->kind().toDisplayString());
auto block = linear->addBlock();
WithInsertPoint guard(block);
auto results = GradientHelper(node).gradient(grad_values);
return fmap(results, [block, linear](Value* grad) -> Value* {
if (!grad)
return nullptr;
block->registerOutput(grad);
return linear->addOutput()->copyMetadata(grad);
});
}
struct ReverseDetails {
ReverseDetails(value_map&& grad_map, Block* reverse_block)
: grad_map(std::move(grad_map)), reverse_block(reverse_block) {}
value_map grad_map;
Block* reverse_block;
};
// AutogradAdd is a special addition function that handles Undef
// AutogradAdd(a, b) == a + b if defined(a) and defined(b)
// AutogradAdd(Undef, b) == b
// AutogradAdd(a, Undef) == a
// AutogradAdd(Undef, Undef) == Undef
static Value* createAutogradAdd(Value* a, Value* b) {
auto graph = a->owningGraph();
return graph->insertNode(graph->create(prim::AutogradAdd, {a, b}))->output();
}
// Before:
// - grad_desc has field f initialized to the original 0-stage graph
// After:
// - the last node of f (f->nodes().reverse()[0]) is a gradient node
// whose block has vjp inputs for all outputs that require_grad
// and vjp outputs for all primal inputs that require_grad
// - grad_desc has df_input_vjps and df_output_vjps set
// (but df_input_vjps will be modified later as well)
static ReverseDetails addReverseInline(Gradient& grad_desc) {
auto& graph = *grad_desc.f;
// note: reverse_node is intentionally not inserted to avoid
// accidentally acting on it (e.g. in elminate dead code),
// std::cout << *reverse_node << to view its state.
auto reverse_node = graph.create(prim::Reverse, 0);
auto reverse_block = reverse_node->addBlock();
WithInsertPoint guard(reverse_block);
value_map grad_map; // x -> dx mapping
const auto get_grad = [&](Value* v) -> Value* {
auto it = grad_map.find(v);
if (it == grad_map.end()) {
auto undef = graph.insertNode(graph.createUndefined());
std::tie(it, std::ignore) = grad_map.emplace(v, undef->output());
}
return it->second;
};
const auto set_grad = [&](Value* x, Value* dx) {
if (Value* prev_grad = grad_map[x]) {
grad_map[x] = createAutogradAdd(prev_grad, dx);
} else {
grad_map[x] = dx;
}
};
auto outputs = graph.outputs();
for (size_t i = 0, num_outputs = outputs.size(); i < num_outputs; ++i) {
Value* output = outputs[i];
if (!output->requires_grad())
continue;
Value* output_grad = reverse_block->addInput()->setType(output->type());
set_grad(output, output_grad);
grad_desc.df_input_vjps.push_back(i);
}
for (auto it = graph.nodes().rbegin(), end = graph.nodes().rend(); it != end;
++it) {
Node* node = *it;
auto inputs = node->inputs();
auto outputs = node->outputs();
if (std::all_of(outputs.begin(), outputs.end(), [](Value* v) {
return !v->requires_grad();
})) {
continue;
}
value_list grad_inputs =
linearGradientForNode(node, fmap(node->outputs(), get_grad));
LowerSimpleTuples(reverse_block);
AT_ASSERT(grad_inputs.size() == node->inputs().size());
for (size_t i = 0, num_inputs = grad_inputs.size(); i < num_inputs; ++i) {
if (!inputs[i]->requires_grad())
continue;
// NB: Not returning a gradient w.r.t. a value that requires grad is
// normal if the input is non-differentiable. This happens e.g. in the
// aten::type_as case.
if (!grad_inputs[i])
continue;
set_grad(inputs[i], grad_inputs[i]);
}
}
auto inputs = graph.inputs();
for (size_t i = 0, num_inputs = inputs.size(); i < num_inputs; ++i) {
Value* input = inputs[i];
if (!input->requires_grad())
continue;
// NB: Not having a gradient defined w.r.t. an input to the graph which
// requires grad can happen and is not an error. It might have been used
// only in non-differentiable contexts (e.g. as second input to
// aten::type_as). In that case we simply ignore it as an output, because it
// won't ever produce any meaningful values.
if (grad_map.count(input) == 0)
continue;
reverse_block->registerOutput(get_grad(input));
grad_desc.df_output_vjps.push_back(i);
}
return ReverseDetails(std::move(grad_map), reverse_block);
}
// Returns a topologically-sorted list of values produced in f, and used in its
// reverse program.
static value_list getReverseCaptures(Gradient& grad_desc) {
auto& graph = *grad_desc.f;
auto primal_block = graph.block();
value_set reverse_captures_set;
value_list reverse_captures; // Invariant: topo sorted
auto check_uses = [&](Value* v) {
for (auto use : v->uses()) {
if (use.user->owningBlock() == primal_block)
continue;
if (/* bool unseen = */ reverse_captures_set.emplace(v).second) {
reverse_captures.push_back(v);
}
}
};
for (Value* input : graph.inputs()) {
check_uses(input);
}
for (Node* node : graph.nodes()) {
for (Value* output : node->outputs())
check_uses(output);
}
return reverse_captures;
}
// Any temporary value from the primal graphs needs to be captured for later use
// in the reverse graph, to avoid costly recomputations. However, a lot of the
// nodes we have in our graphs are simply constants, which are cheap to execute
// and replicate, and so it's better to just copy them into the reverse graph,
// without polluting the output lists unnecessarily.
static void liftConstants(Gradient& grad_desc, ReverseDetails& rev_info) {
static const auto err = [](Value*) -> Value* {
throw std::runtime_error("unexpected input");
};
auto& graph = *grad_desc.f;
Block* reverse_block = rev_info.reverse_block;
for (Node* top_node : reverse_block->nodes()) {
AT_ASSERT(
top_node->kind() == prim::GradOf ||
top_node->kind() == prim::AutogradAdd ||
top_node->kind() == prim::Undefined);
if (top_node->kind() != prim::GradOf)
continue;
Block* grad_body = top_node->blocks().at(0);
for (Node* node : grad_body->nodes()) {
for (Value* input : node->inputs()) {
if (input->node()->kind() != prim::Constant)
continue;
if (input->node()->owningBlock() == grad_body)
continue;
Node* lifted_constant = graph.createClone(input->node(), err);
reverse_block->prependNode(lifted_constant);
node->replaceInputWith(input, lifted_constant->output());
}
}
}
}
static void deduplicateSizeCaptures(
Gradient& grad_desc,
ReverseDetails& rev_info) {
Block* primal_block = grad_desc.f->block();
const auto usedOnlyInReverse = [primal_block](Value* v) {
const auto& uses = v->uses();
return std::all_of(uses.begin(), uses.end(), [primal_block](const Use& u) {
return u.user->owningBlock() != primal_block;
});
};
auto captures = getReverseCaptures(grad_desc);
value_set capture_set(captures.begin(), captures.end());
for (Value* capture : captures) {
Node* node = capture->node();
if (!node->matches("aten::size(Tensor self) -> int[]")) {
continue;
}
if (usedOnlyInReverse(capture) && capture_set.count(node->input())) {
WithInsertPoint insert_guard{*rev_info.reverse_block->nodes().begin()};
capture->replaceAllUsesWith(SymbolicVariable(node->input()).size());
node->destroy();
}
}
}
static void eliminateDeadCode(ReverseDetails& rev_info) {
// addReverseInline has to call gradientForNode if *any* of the inputs
// require grad, but it will emit vjps for *all* inputs. Use DCE to remove
// unnecessary nodes. Additionally, requires_grad() on intermediates is an
// overapproximation of the real state, so we might have emitted some
// gradients, only to realize that they were unnecessary once we reach a
// point that doesn't require grad.
// Of course, we need to filter out corresponding entries of grad_map, because
// we don't want to accidentally access freed pointers later.
std::function<void(const std::unordered_set<const Value*>&)> cb =
[&](const std::unordered_set<const Value*>& live_values) {
std::vector<Value*> to_erase;
for (auto& entry : rev_info.grad_map) {
if (!live_values.count(entry.second)) {
to_erase.push_back(entry.first);
}
}
for (Value* v : to_erase) {
rev_info.grad_map.erase(v);
}
};
EliminateDeadCode(rev_info.reverse_block, std::move(cb));
}
static void Optimize(Gradient& grad_desc, ReverseDetails& rev_info) {
// TODO: we are sometimes emitting expressions like SumToSize(SumToSize(x,
// s1), s2), which are equivalent to SumToSize(x, s2), and could save us some
// captures, but I'm not 100% sure how to optimize this at this stage, since
// we don't know which GradOf blocks will be stitched together to form the
// derivative. I guess a smart analysis could implement this, but I didn't
// have time before the 1.0 release, so I put this only as a peephole
// optimization.
liftConstants(grad_desc, rev_info);
// We generally add a lot of aten::size calls (for derivatives of broadcasting
// operators), and they often end up duplicated, and would get captured
// multiple times. Make sure we deduplicate them before lifting.
EliminateCommonSubexpression(grad_desc.f);
deduplicateSizeCaptures(grad_desc, rev_info);
eliminateDeadCode(rev_info);
}
// Takes a grad_desc.f returned from `addReverseInline` and splits off the
// reverse_block into its own graph, storing it in df.
// All intermediates needed in the second stage are added to
// outputs of f, and taken as inputs in df. For a more
// detailed description see Note [Gradient graphs] in autodiff.h.
// This function also initializes the fields in grad_desc that were undefined
// after `addReverseInline` (and extends `df_input_vjps` with vjps for captured
// temporaries).
static void lambdaLiftReverse(Gradient& grad_desc, ReverseDetails& rev_info) {
auto& graph = *grad_desc.f;
auto primal_block = graph.block();
auto reverse_block = rev_info.reverse_block;
// --------------------------------------------------------------------------
// 1. Find values of f that need to be captured.
// --------------------------------------------------------------------------
// First, we need to find all values that are produced in f,
// and used in df. They will need to be added as inputs of the df
// and some of them may also need to be appended as outputs of f if
// they are not already an input or an output of f
// Invariant: topo sorted
value_list reverse_captures = getReverseCaptures(grad_desc);
// --------------------------------------------------------------------------
// 2. Prepare input/outputs lists for f and df
// --------------------------------------------------------------------------
// It's simple to construct primal_inputs/reverse_outputs,
// but primal_outputs/reverse_inputs are much more subtle.
// Here's a summary of how they are supposed to look like:
//
// Primal outputs:
// [original outputs], [temporaries]
//
// Reverse inputs:
// [output vjps (aka grad_outputs)], [temporary vjps]
// [captured primal values, in topological order],
// -- Construct primal_outputs, df_input_captures, f_real_outputs ----
grad_desc.f_real_outputs = graph.outputs().size();
std::unordered_map<Value*, size_t> orig_primal_outputs_idx;
std::unordered_map<Value*, size_t> orig_primal_inputs_idx;
// NOTE: we use emplace to avoid replacing an existing index if an output is
// repeated
for (size_t i = 0, num_outputs = graph.outputs().size(); i < num_outputs; ++i)
orig_primal_outputs_idx.emplace(graph.outputs()[i], i);
for (size_t i = 0, num_inputs = graph.inputs().size(); i < num_inputs; ++i)
orig_primal_inputs_idx[graph.inputs()[i]] = i;
// NB: reverse_captures are already deduplicated, and in topo order
for (Value* capture_val : reverse_captures) {
// If it's already an output we don't have to add anything,
// but register the fact that it needs to be captured.
if (orig_primal_outputs_idx.count(capture_val) > 0) {
grad_desc.df_input_captured_outputs.push_back(
orig_primal_outputs_idx[capture_val]);
// If it's an input, we could add it as an output but in fact it's
// more efficient to use a special kind of capture.
} else if (orig_primal_inputs_idx.count(capture_val) > 0) {
grad_desc.df_input_captured_inputs.push_back(
orig_primal_inputs_idx.at(capture_val));
// Otherwise it's just a regular intermediate value that we need to add as
// an output
} else {
// we need to create a new temporary output for this capture because it
// wasn't availiable.
graph.registerOutput(capture_val);
grad_desc.df_input_captured_outputs.emplace_back(
graph.outputs().size() - 1);
}
}
// -- Add VJPs for temporaries, adjust df_input_vjps -------------------------
// NB [possible optimization]: use the newly added vjp input as soon as the
// first vjp for that value is generated, to reduce the lifespan of this input
// (currently we add it to the final vjp after all adds).
for (size_t i = grad_desc.f_real_outputs; i < graph.outputs().size(); ++i) {
Value* tmp = graph.outputs().at(i);
// Add VJP inputs only for intermediates that actually required grad.
// Note that we check the contents of the grad_map instead of
// tmp->requires_grad(), becuase it's actually a more faithful source.
// tmp->requires_grad() is really an overapproximation (i.e. it can have
// false positives), while the gradients we will emit for this value can get
// DCE-d in the optimization pass (because it has no influence on the real
// f's outputs that we differentiate).
if (rev_info.grad_map.count(tmp) == 0)
continue;
Value* tmp_vjp_in = reverse_block->addInput()->setType(tmp->type());
Value* tmp_vjp_prev = rev_info.grad_map.at(tmp);
// This is quite weird because we can't first make a sum and then replace
// all uses of tmp_vjp_prev (that would replace its use in the sum too!), so
// we create an incorrect sum that doesn't use prev vjp, replace uses, and
// fix the sum.
Value* new_vjp = createAutogradAdd(tmp_vjp_in, tmp_vjp_in);
new_vjp->node()->moveAfter(tmp_vjp_prev->node());
tmp_vjp_prev->replaceAllUsesWith(new_vjp);
new_vjp->node()->replaceInput(1, tmp_vjp_prev);
grad_desc.df_input_vjps.emplace_back(i);
}
// add the captures as formal arguments to the reverse_block
// afterward inputs: [output vjps][temporary vjps][captures]
// construct a map from captured 'value' to the index in the input list
// used to extract this block into its own function
std::unordered_map<Value*, size_t> capture_to_formal_index;
const auto& add_capture = [&](Value* captured) {
capture_to_formal_index[captured] = reverse_block->inputs().size();
reverse_block->addInput()->copyMetadata(captured);
};
for (auto& offset : grad_desc.df_input_captured_inputs)
add_capture(graph.inputs()[offset]);
for (auto& offset : grad_desc.df_input_captured_outputs)
add_capture(graph.outputs()[offset]);
grad_desc.df = std::make_shared<Graph>();
grad_desc.df->block()->cloneFrom(reverse_block, [&](Value* v) {
return grad_desc.df->inputs()[capture_to_formal_index.at(v)];
});
// reverse_node was just to hold onto reverse_block in a debuggable way
// we can remove it now.
reverse_block->owningNode()->destroy();
}
Gradient differentiate(std::shared_ptr<Graph>& graph) {
Gradient grad_desc;
// Take ownership of the graph
AT_CHECK(
graph.use_count() == 1,
"differentiate will mutate and destroy the graph, so it requires "
"graph.use_count() == 1, but found %d",
graph.use_count());
std::swap(graph, grad_desc.f);
// XXX: Take care when handling outputs - they can be duplicated!
WithInsertPoint guard(grad_desc.f->block());
// Fills in df_input_vjps and df_output_vjps
auto rev_info = addReverseInline(grad_desc);
Optimize(grad_desc, rev_info);
// Clean up old nodes which has been replaced by forward graphs in torchscript
EliminateDeadCode(grad_desc.f->block());
// Fills in f, df, f_real_outputs, df_input_captures,
// modifies df_input_vjps (new vjps are added for temporaries)
lambdaLiftReverse(grad_desc, rev_info);
// It's possible the we've cloned the same constants many times, so
// de-duplicate them
ConstantPooling(grad_desc.df);
return grad_desc;
}
} // namespace jit
} // namespace torch