pytorch自定义二值化网络层方式

yipeiwu_com6年前Python基础

任务要求:

自定义一个层主要是定义该层的实现函数,只需要重载Function的forward和backward函数即可,如下:

import torch
from torch.autograd import Function
from torch.autograd import Variable

定义二值化函数

class BinarizedF(Function):
  def forward(self, input):
    self.save_for_backward(input)
    a = torch.ones_like(input)
    b = -torch.ones_like(input)
    output = torch.where(input>=0,a,b)
    return output
  def backward(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)
    return input_grad

定义一个module

class BinarizedModule(nn.Module):
  def __init__(self):
    super(BinarizedModule, self).__init__()
    self.BF = BinarizedF()
  def forward(self,input):
    print(input.shape)
    output =self.BF(input)
    return output

进行测试

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

class BinarizedF(Function):
  def forward(self, input):
    self.save_for_backward(input)
    output = torch.ones_like(input)
    output[input<0] = -1
    return output
  def backward(self, output_grad):
    input, = self.saved_tensors
    input_grad = output_grad.clone()
    input_abs = torch.abs(input)
    input_grad[input_abs>1] = 0
    return input_grad

以上这篇pytorch自定义二值化网络层方式就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

pyqt5之将textBrowser的内容写入txt文档的方法

如下所示: try: StrText = self.textBrowser.toPlainText() qS = str(StrText)...

Python实现的计数排序算法示例

Python实现的计数排序算法示例

本文实例讲述了Python实现的计数排序算法。分享给大家供大家参考,具体如下: 计数排序是一种非常快捷的稳定性强的排序方法,时间复杂度O(n+k),其中n为要排序的数的个数,k为要排序的...

PyTorch 解决Dataset和Dataloader遇到的问题

今天在使用PyTorch中Dataset遇到了一个问题。先看代码 class psDataset(Dataset): def __init__(self, x, y, trans...

python使用selenium实现批量文件下载

python使用selenium实现批量文件下载

背景 实现需求:批量下载联想某型号的全部驱动程序。 一般在做网络爬虫的时候,都是保存网页信息为主,或者下载单个文件。当涉及到多文件批量下载的时候,由于下载所需时间不定,下载的文件名不定,...

Pandas的read_csv函数参数分析详解

函数原型 复制代码 代码如下:pd.read_csv(filepath_or_buffer, sep=',', delimiter=None, header='infer', name...