René's URL Explorer Experiment

    {
       "@context": "https://schema.org",
       "@type": "Article",
       "name": "Custom Autograd Functions",
       "headline": "Custom Autograd Functions",
       "description": "Custom autograd functions in PyTorch C++ \u2014 defining forward and backward passes with torch::autograd::Function.",
       "url": "/api/autograd/custom_functions.html",
       "articleBody": "Custom Autograd Functions# PyTorch allows you to define custom autograd functions with custom forward and backward implementations. Function Base Class# template\u003cclass T\u003estruct Function# To use custom autograd operations, implement a Function subclass with static forward and backward functions: forward can take as many arguments as you want and should return either a variable list or a Variable. Use of any direct Variable arguments will be registered in the graph but no vectors/sets or any other data structures will be traversed. You can use std::optional\u003cTensor\u003e as one of the arguments and it will be registered as a variable in the graph if the argument has a value. It should take a pointer to torch::autograd::AutogradContext as the first argument. Variables can be saved in the ctx using ctx-\u003esave_for_backward (see torch::autograd::AutogradContext::save_for_backward) and other data can be saved in the ctx-\u003esaved_data map (see torch::autograd::AutogradContext::saved_data) in the form of \u003cstd::string, at::IValue\u003e pairs. backward should take a pointer to torch::autograd::AutogradContext and a variable list containing as many Variables as there were outputs from forward as arguments. It should return as many Variables as there were inputs with each of them containing the gradient w.r.t. its corresponding input. Variables saved in forward can be accessed with ctx-\u003eget_saved_variables (see torch::autograd::AutogradContext::get_saved_variables) and other saved data can be accessed from ctx-\u003esaved_data. To enable compiled autograd support (torch.compile for backward) for your custom autograd operation, you can set MyFunction::is_traceable (see Function::istraceable notes below). For example: class MyFunction : public Function\u003cMyFunction\u003e { public: static constexpr bool is_traceable = true; static variable_list forward(AutogradContext *ctx, int n, Variable var) { // Save data for backward in context ctx-\u003esaved_data[\"n\"] = n; var.mul_(n); // Mark var as modified by inplace operation ctx-\u003emark_dirty({var}); return {var}; } static variable_list backward(AutogradContext *ctx, variable_list grad_output) { // Use data saved in forward auto n = ctx-\u003esaved_data[\"n\"].toInt(); return {grad_output[0]*n}; } }; To use MyFunction: Variable x; auto y = MyFunction::apply(6, x); // Example backward call y[0].sum().backward(); Public Static Functions template\u003ctypename X = T, typename ...Args\u003estatic auto apply(Args\u0026\u0026... args) -\u003e std::enable_if_t\u003cstd::is_same_v\u003cX, T\u003e, forward_t\u003cX, Args...\u003e\u003e# Public Static Attributes static constexpr bool is_traceable = false# AutogradContext# struct AutogradContext# Context to save information during forward that can be accessed in backward in custom autograd operations (see torch::autograd::Function for details). Public Functions AutogradContext() = default# AutogradContext(const AutogradContext \u0026other) = delete# AutogradContext \u0026operator=(const AutogradContext \u0026other) = delete# AutogradContext(AutogradContext \u0026\u0026other) = delete# AutogradContext \u0026operator=(AutogradContext \u0026\u0026other) = delete# ~AutogradContext() = default# AutogradContext(PackedArgs \u0026packed_args)# void save_for_backward(variable_list to_save)# Saves the list of variables for a future call to backward. This should be called at most once from inside of forward. void mark_dirty(const variable_list \u0026inputs)# Marks variables in the list as modified in an in-place operation. This should be called at most once from inside of forward and all arguments should be inputs. void mark_non_differentiable(const variable_list \u0026outputs)# Marks outputs in the list as not requiring gradients. This should be called at most once from inside of forward and all arguments should be outputs. void set_materialize_grads(bool value)# variable_list get_saved_variables() const# Get the list of variables that were saved in forward using save_for_backward(). Before returning them to the user, a check is made to ensure that they were not modified by any in-place operations. const std::unordered_set\u003cat::TensorImpl*\u003e \u0026get_and_bump_dirty() const# const std::unordered_set\u003cat::TensorImpl*\u003e \u0026get_non_differentiable() const# bool needs_input_grad(size_t output_edge_index) const# Expose the Node\u2019s task_should_compute_output method to the cpp custom autograd Function as needs_input_grad. bool needs_input_grad(std::initializer_list\u003cIndexRange\u003e idxs) const# Public Members ska::flat_hash_map\u003cstd::string, at::IValue\u003e saved_data# Can be used to save non-variable data for backward. Creating Custom Functions# To create a custom autograd function, inherit from torch::autograd::Function and implement the static forward and backward methods: Example: class MyReLU : public torch::autograd::Function\u003cMyReLU\u003e { public: static torch::Tensor forward( torch::autograd::AutogradContext* ctx, torch::Tensor input) { ctx-\u003esave_for_backward({input}); return input.clamp_min(0); } static torch::autograd::variable_list backward( torch::autograd::AutogradContext* ctx, torch::autograd::variable_list grad_outputs) { auto saved = ctx-\u003eget_saved_variables(); auto input = saved[0]; auto grad_output = grad_outputs[0]; auto grad_input = grad_output * (input \u003e 0).to(grad_output.dtype()); return {grad_input}; } }; // Usage auto output = MyReLU::apply(input); Custom Kernels and AutoDispatchBelowADInplaceOrView# For users implementing custom kernels who want to redispatch below Autograd dispatch keys, use at::AutoDispatchBelowADInplaceOrView instead of InferenceMode: class ROIAlignFunction : public torch::autograd::Function\u003cROIAlignFunction\u003e { public: static torch::autograd::variable_list forward( torch::autograd::AutogradContext* ctx, const torch::autograd::Variable\u0026 input, const torch::autograd::Variable\u0026 rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio, bool aligned) { ctx-\u003esaved_data[\"spatial_scale\"] = spatial_scale; ctx-\u003esaved_data[\"pooled_height\"] = pooled_height; ctx-\u003esaved_data[\"pooled_width\"] = pooled_width; ctx-\u003esaved_data[\"sampling_ratio\"] = sampling_ratio; ctx-\u003esaved_data[\"aligned\"] = aligned; ctx-\u003esaved_data[\"input_shape\"] = input.sizes(); ctx-\u003esave_for_backward({rois}); at::AutoDispatchBelowADInplaceOrView guard; auto result = roi_align( input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned); return {result}; } }; For customized inplace and view kernels, see the custom kernel tutorial for more details.",
       "author": {
         "@type": "Organization",
         "name": "PyTorch Contributors",
         "url": "https://pytorch.org"
       },
       "image": "https://pytorch.org/docs/stable/_static/img/pytorch_seo.png",
       "mainEntityOfPage": {
         "@type": "WebPage",
         "@id": "/api/autograd/custom_functions.html"
       },
       "datePublished": "2023-01-01T00:00:00Z",
       "dateModified": "2023-01-01T00:00:00Z"
     }