其实各大深度学习框架背后的原理都可以理解为拟合一个参数数量特别庞大的函数，所以各框架都能用来拟合任意函数，Pytorch也能。在这篇博客中，就以拟合y = ax + b为例（a和b为需要拟合的参数），说明在Pytorch中如何拟合一个函数。

一、定义拟合网络

1、观察普通的神经网络的优化流程

# 定义网络
net = ...
# 定义优化器
optimizer = torch.optim.Adam(net.parameters(), lr=0.001, weight_decay=0.0005)
# 定义损失函数
loss_op = torch.nn.MSELoss(reduction='sum')
# 优化
for step, (inputs, tag) in enumerate(dataset_loader):
	# 向前传播
	outputs = net(inputs)
	# 计算损失
	loss = loss_op(tag, outputs)
	# 清空梯度
	optimizer.zero_grad()
	# 向后传播
	loss.backward()
	# 更新梯度
	optimizer.step()

上面的代码就是一般情况下的流程。为了能使用Pytorch内置的优化器，所以我们需要定义一个一个网络，实现函数parameters（返回需要优化的参数）和forward（向前传播）；为了能支持GPU优化，还需要实现cuda和cpu两个函数，把参数从内存复制到GPU上和从GPU复制回内存。基于以上要求，网络的定义就类似于：

class Net:
    def __init__(self):
        # 在这里定义要求的参数
        pass

    def cuda(self):
        # 传输参数到GPU
        pass

    def cpu(self):
        # 把参数传输回内存
        pass

    def forward(self, inputs):
    	# 实现向前传播，就是根据输入inputs计算一遍输出
        pass

    def parameters(self):
    	# 返回参数
        pass

在拟合数据量很大时，还可以使用GPU来加速；如果没有英伟达显卡，则可以不实现cuda和cpu这两个函数。

2、初始化网络

回顾本文目的，拟合： y = ax + b, 所以在__init__函数中就需要定义a和b两个参数，另外为了实现parameters、cpu和cuda，还需要定义属性__parameters和__gpu：

    def __init__(self):
        # y = a*x + b
        self.a = torch.rand(1, requires_grad=True)  # 参数a
        self.b = torch.rand(1, requires_grad=True)  # 参数b
        self.__parameters = dict(a=self.a, b=self.b)  # 参数字典
        self.___gpu = False  # 是否使用gpu来拟合

要拟合的参数，不能初始化为0！ ，一般使用随机值即可。还需要把requires_grad参数设置为True，这是为了支持向后传播。

3、实现向前传播

    def forward(self, inputs):
        return self.a * inputs + self.b

非常的简单，就是根据输入inputs计算一遍输出，在本例中，就是计算一下 y = ax + b。计算完了要记得返回计算的结果。

4、把参数传送到GPU

为了支持GPU来加速拟合，需要把参数传输到GPU，且需要更新参数字典__parameters：

    def cuda(self):
        if not self.___gpu:
            self.a = self.a.cuda().detach().requires_grad_(True)  # 把a传输到gpu
            self.b = self.b.cuda().detach().requires_grad_(True)  # 把b传输到gpu
            self.__parameters = dict(a=self.a, b=self.b)  # 更新参数
            self.___gpu = True  # 更新标志，表示参数已经传输到gpu了
        # 返回self，以支持链式调用
        return self

参数a和b，都是先调用detach再调用requires_grad_，是为了避免错误raise ValueError("can't optimize a non-leaf Tensor")（参考：ValueError: can’t optimize a non-leaf Tensor?）。

4、把参数传输回内存

类似于cuda函数，不做过多解释。

    def cpu(self):
        if self.___gpu:
            self.a = self.a.cpu().detach().requires_grad_(True)
            self.b = self.b.cpu().detach().requires_grad_(True)
            self.__parameters = dict(a=self.a, b=self.b)
            self.___gpu = False
        return self

5、返回网络参数

为了能使用Pytorch内置的优化器，就要实现parameters函数，观察Pytorch里面的实现：

    def parameters(self, recurse=True):
        r"""... """
        for name, param in self.named_parameters(recurse=recurse):
            yield param

实际上就是使用yield返回网络的所有参数，因此本例中的实现如下：

    def parameters(self):
        for name, param in self.__parameters.items():
            yield param

完整的实现将会放在后面。

二、测试

1、生成测试数据

def main():
    # 生成虚假数据
    x = np.linspace(1, 50, 50)
    # 系数a、b
    a = 2
    b = 1
    # 生成y
    y = a * x + b
    # 转换为Tensor
    x = torch.from_numpy(x.astype(np.float32))
    y = torch.from_numpy(y.astype(np.float32))

2、定义网络

	# 定义网络
    net = Net()
    # 定义优化器
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001, weight_decay=0.0005)
    # 定义损失函数
    loss_op = torch.nn.MSELoss(reduction='sum')

3、把数据传输到GPU（可选）

	# 传输到GPU
    if torch.cuda.is_available():
        x = x.cuda()
        y = y.cuda()
        net = net.cuda()

4、定义优化器和损失函数

如果要使用GPU加速，优化器必须要在网络的参数传输到GPU之后在定义，否则优化器里的参数还是内存里的那些参数，传到GPU里面的参数不能被更新。 可以根据代码来理解这句话。

	# 定义优化器
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001, weight_decay=0.0005)
    # 定义损失函数
    loss_op = torch.nn.MSELoss(reduction='sum')

5、拟合（也是优化）

	# 最多优化20001次
    for i in range(1, 20001, 1):
    	# 向前传播
        out = net.forward(x)
		# 计算损失
        loss = loss_op(y, out)
		# 清空梯度（非常重要）
        optimizer.zero_grad()
		# 向后传播，计算梯度
        loss.backward()
		# 更新参数
        optimizer.step()
		# 得到损失的numpy值
        loss_numpy = loss.cpu().detach().numpy()
        if i % 1000 == 0:  # 每1000次打印一下损失
            print(i, loss_numpy)

        if loss_numpy < 0.00001:  # 如果损失小于0.00001
        	# 打印参数
        	a = net.a.cpu().detach().numpy()
        	b = net.b.cpu().detach().numpy()
            print(a, b)
            # 退出
            exit()

6、完整示例代码

# coding=utf-8
from __future__ import absolute_import, division, print_function
import torch
import numpy as np


class Net:
    def __init__(self):
        # y = a*x + b
        self.a = torch.rand(1, requires_grad=True)  # 参数a
        self.b = torch.rand(1, requires_grad=True)  # 参数b
        self.__parameters = dict(a=self.a, b=self.b)  # 参数字典
        self.___gpu = False  # 是否使用gpu来拟合

    def cuda(self):
        if not self.___gpu:
            self.a = self.a.cuda().detach().requires_grad_(True)  # 把a传输到gpu
            self.b = self.b.cuda().detach().requires_grad_(True)  # 把b传输到gpu
            self.__parameters = dict(a=self.a, b=self.b)  # 更新参数
            self.___gpu = True  # 更新标志，表示参数已经传输到gpu了
        # 返回self，以支持链式调用
        return self

    def cpu(self):
        if self.___gpu:
            self.a = self.a.cpu().detach().requires_grad_(True)
            self.b = self.b.cpu().detach().requires_grad_(True)
            self.__parameters = dict(a=self.a, b=self.b)  # 更新参数
            self.___gpu = False
        return self

    def forward(self, inputs):
        return self.a * inputs + self.b

    def parameters(self):
        for name, param in self.__parameters.items():
            yield param


def main():

    # 生成虚假数据
    x = np.linspace(1, 50, 50)

    # 系数a、b
    a = 2
    b = 1

    # 生成y
    y = a * x + b

    # 转换为Tensor
    x = torch.from_numpy(x.astype(np.float32))
    y = torch.from_numpy(y.astype(np.float32))

    # 定义网络
    net = Net()

    # 传输到GPU
    if torch.cuda.is_available():
        x = x.cuda()
        y = y.cuda()
        net = net.cuda()

    # 定义优化器
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001, weight_decay=0.0005)

    # 定义损失函数
    loss_op = torch.nn.MSELoss(reduction='sum')

    # 最多优化20001次
    for i in range(1, 20001, 1):
        # 向前传播
        out = net.forward(x)
        # 计算损失
        loss = loss_op(y, out)
        # 清空梯度（非常重要）
        optimizer.zero_grad()
        # 向后传播，计算梯度
        loss.backward()
        # 更新参数
        optimizer.step()
        # 得到损失的numpy值
        loss_numpy = loss.cpu().detach().numpy()
        if i % 1000 == 0:  # 每1000次打印一下损失
            print(i, loss_numpy)

        if loss_numpy < 0.00001:  # 如果损失小于0.00001
            # 打印参数
            a = net.a.cpu().detach().numpy()
            b = net.b.cpu().detach().numpy()
            print(a, b)
            # 退出
            exit()


if __name__ == '__main__':
    main()