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