pytorch-design NEW Function and Module

Design NEW Function and Module

class LinearFunc(torch.autograd.Function):

@staticmethod
def forward(ctx, input, weight, bias=None):
ctx.save_for_backward(input, weight, bias)
output = input.mm(weight)
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
pass
return output

@staticmethod
def backward(ctx, grad_output):
input, weight, bias = ctx.saved_variables
return None, None, None

class Linear(nn.Module):
def __init__(self, xd, yd, bias=True):
super(Linear, self).__init__()
self.xd = xd
self.yd = yd
self.weight = nn.Parameter( torch.FloatTensor(xd, yd) )
if bias:
self.bias = nn.Parameter( torch.FloatTensor(yd) )
else:
self.register_parameter('bias', None)
self.weight.data.uniform_(-0.1, 0.1)
if bias is not None:
self.bias.data.zero_()
pass

def forward(self, x):
return LinearFunc.apply(x, self.weight, self.bias)

-------------------------

0. http://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html#

1. https://www.qcloud.com/community/article/831497

2. https://discuss.pytorch.org/t/difference-of-methods-between-torch-nn-and-functional/1076

3. https://discuss.pytorch.org/t/understanding-how-torch-nn-module-works/122

4. https://discuss.pytorch.org/t/how-to-choose-between-torch-nn-functional-and-torch-nn-module-see-mnist-https-github-com-pytorch-examples-blob-master-mnist-main-py/2800/10