Pytorch使用autograd.function自定义op
2021/8/17 6:06:48
本文主要是介绍Pytorch使用autograd.function自定义op,对大家解决编程问题具有一定的参考价值,需要的程序猿们随着小编来一起学习吧!
0x01 简介
Pytorch是利用Variable与Function来构建计算图的,Function就像计算图中的边,实现Variable的计算,并输出新的Variable,因此Function需要有forward和backward的功能,前者用于正常传递输入,后者用于计算梯度.(但直接用nn.Module构建就可以只用写forward,backward调用自动求导计算了,另用Module可以保存参数而Function不能,因此前者多用于写一些需要保存参数的如自定义的层,而后者通常用来写一个操作,如激活函数之类的,偷懒直接nn.Module就行了).
0x02 使用
pytorch官方文档上写的是得用save_for_backward保存下输入,这个保存的数据会在backward的时候通过saved_tensors读取
使用时调用apply方法即可。
或者直接重命名一下
你已经学会了,现在来试试手写一个Relu吧.jpg
import torch from torch.autograd import Variable from torch.autograd import Function class Myop1(Function): @staticmethod def forward(self, input): self.save_for_backward(input) output = input.clamp(min=0) return output @staticmethod def backward(self, grad_output): ## input = self.saved_tensors grad_input = grad_output.clone() grad_input[input<0] = 0 return grad_input
看一下Variable与Function的关系
input_1 = Variable(torch.randn(1), requires_grad=True) print(input_1) relu2 = Myop1.apply output_= relu2(input_1) print(output_.grad_fn)
输出
<torch.autograd.function.Myop1Backward object at 0x00000123A49A29E0>
or封装成函数
def relu(input): return Myop1.apply(input)
这篇关于Pytorch使用autograd.function自定义op的文章就介绍到这儿,希望我们推荐的文章对大家有所帮助,也希望大家多多支持为之网!
- 2024-05-29Elasticsearch慢查询日志配置
- 2024-05-29揭秘华为如此多成功项目的产品关键——Charter模板
- 2024-05-29海外IDC业务拓展的7大挑战
- 2024-05-29InLine Chat功能优化对标Github Copilot,CodeGeeX带来更高效、更直观的编程体验!
- 2024-05-29CodeGeeX 智能编程助手 6 项功能升级,在Visual Studio插件市场霸榜2周!
- 2024-05-29AutoMQ 生态集成 Apache Doris
- 2024-05-292024年IDC行业的深度挖掘:机遇、挑战与未来展望
- 2024-05-29五款扩展组件齐发 —— Volcano、Keda、Crane-scheduler 等,邀你体验
- 2024-05-29AutoMQ 对象存储数据高效组织的秘密: Compaction
- 2024-05-29活动预告|来 GIAC 大会听大数据降本利器:AutoMQ 基于云原生重新设计的 Kafka