python二值化函数_pytorch自定义二值化网络层方式

任务要求：

自定义一个层主要是定义该层的实现函数,只需要重载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(