pytorch自定义二值化网络层方式
任务要求：
自定义一个层主要是定义该层的实现函数,只需要重载Function的forward和backward函数即可,如下:
importtorch fromtorch.autogradimportFunction fromtorch.autogradimportVariable
定义二值化函数
classBinarizedF(Function): defforward(self,input): self.save_for_backward(input) a=torch.ones_like(input) b=-torch.ones_like(input) output=torch.where(input>=0,a,b) returnoutput defbackward(self,output_grad): input,=self.saved_tensors input_abs=torch.abs(input) ones=torch.ones_like(input) zeros=torch.zeros_like(input) input_grad=torch.where(input_abs<=1,ones,zeros) returninput_grad
定义一个module
classBinarizedModule(nn.Module): def__init__(self): super(BinarizedModule,self).__init__() self.BF=BinarizedF() defforward(self,input): print(input.shape) output=self.BF(input) returnoutput
进行测试
a=Variable(torch.randn(4,480,640),requires_grad=True) output=BinarizedModule()(a) output.backward(torch.ones(a.size())) print(a) print(a.grad)
其中,二值化函数部分也可以按照方式写,但是速度慢了0.05s
classBinarizedF(Function): defforward(self,input): self.save_for_backward(input) output=torch.ones_like(input) output[input<0]=-1 returnoutput defbackward(self,output_grad): input,=self.saved_tensors input_grad=output_grad.clone() input_abs=torch.abs(input) input_grad[input_abs>1]=0 returninput_grad
以上这篇pytorch自定义二值化网络层方式就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。
声明：本文内容来源于网络，版权归原作者所有，内容由互联网用户自发贡献自行上传，本网站不拥有所有权，未作人工编辑处理，也不承担相关法律责任。如果您发现有涉嫌版权的内容，欢迎发送邮件至：czq8825#qq.com（发邮件时，请将#更换为@）进行举报，并提供相关证据，一经查实，本站将立刻删除涉嫌侵权内容。