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