嘗試使用了pytorch，相比其他深度學(xué)習(xí)框架，pytorch顯得簡(jiǎn)潔易懂。花時(shí)間讀了部分源碼，主要結(jié)合簡(jiǎn)單例子帶著問(wèn)題閱讀，不涉及源碼中C拓展庫(kù)的實(shí)現(xiàn)。

一個(gè)簡(jiǎn)單例子

實(shí)現(xiàn)單層softmax二分類(lèi)，輸入特征維度為4，輸出為2，經(jīng)過(guò)softmax函數(shù)得出輸入的類(lèi)別概率。代碼示意：定義網(wǎng)絡(luò)結(jié)構(gòu)；使用SGD優(yōu)化；迭代一次，隨機(jī)初始化三個(gè)樣例，每個(gè)樣例四維特征，target分別為1,0,1；前向傳播，使用交叉熵計(jì)算loss；反向傳播，最后由優(yōu)化算法更新權(quán)重，完成一次迭代。

1. import torch
2. import torch.nn as nn
3. import torch.nn.functional as F
5. class Net(nn.Module):
7. def __init__(self):
8. super(Net, self).__init__()
9. self.linear = nn.Linear(4, 2)
11. def forward(self, input):
12. out = F.softmax(self.linear(input))
13. return out
15. net = Net()
16. sgd = torch.optim.SGD(net.parameters(), lr=0.001)
17. for epoch in range(1):
18. features = torch.autograd.Variable(torch.randn(3, 4), requires_grad=True)
19. target = torch.autograd.Variable(torch.LongTensor([1, 0, 1]))
20. sgd.zero_grad()
22. out = net(features)
23. loss = F.cross_entropy(out, target)
24. loss.backward()
25. sgd.step()

從上面的例子，帶著下面的問(wèn)題閱讀源碼：

• pytorch的主要概念：Tensor、autograd、Variable、Function、Parameter、Module（Layers）、Optimizer；
• 自定義Module如何組織網(wǎng)絡(luò)結(jié)構(gòu)和網(wǎng)絡(luò)參數(shù)；
• 前向傳播、反向傳播實(shí)現(xiàn)流程
• 優(yōu)化算法類(lèi)如何實(shí)現(xiàn)，如何和自定義Module聯(lián)系并更新參數(shù)。

pytorch的主要概念

pytorch的主要概念官網(wǎng)有很人性化的教程Deep Learning with PyTorch: A 60 Minute Blitz， 這里簡(jiǎn)單概括這些概念：

Tensor

類(lèi)似numpy的ndarrays，強(qiáng)化了可進(jìn)行GPU計(jì)算的特性，由C拓展模塊實(shí)現(xiàn)。如上面的torch.randn(3, 4) 返回一個(gè)3*4的Tensor。和numpy一樣，也有一系列的Operation，如

1. x = torch.rand(5, 3)
2. y = torch.rand(5, 3)
3. print x + y
4. print torch.add(x, y)
5. print x.add_(y)

Varaiable與autograd

Variable封裝了Tensor，包括了幾乎所有的Tensor可以使用的Operation方法，主要使用在自動(dòng)求導(dǎo)(autograd)，Variable類(lèi)繼承_C._VariableBase，由C拓展類(lèi)定義實(shí)現(xiàn)。
Variable是autograd的計(jì)算單元，Variable通過(guò)Function組織成函數(shù)表達(dá)式（計(jì)算圖）：

• data 為其封裝的tensor值
• grad 為其求導(dǎo)后的值
• creator 為創(chuàng)建該Variable的Function，實(shí)現(xiàn)中g(shù)rad_fn屬性則指向該Function。
  如：
  1. import torch
  2. from torch.autograd import Variable
  3. x = Variable(torch.ones(2, 2), requires_grad=True)
  4. y = x + 2
  5. print y.grad_fn
  6. print "before backward: ", x.grad
  7. y.backward()
  8. print "after backward: ", x.grad

  輸出結(jié)果：

  1. <torch.autograd.function.AddConstantBackward object at 0x7faa6f3bdd68>
  2. before backward: None
  3. after backward: Variable containing:
  4. 1
  5. [torch.FloatTensor of size 1x1]

  調(diào)用y的backward方法，則會(huì)對(duì)創(chuàng)建y的Function計(jì)算圖中所有requires_grad=True的Variable求導(dǎo)（這里的x）。例子中顯然dy/dx = 1。

Parameter

? ?Parameter 為Variable的一個(gè)子類(lèi)，后面還會(huì)涉及，大概兩點(diǎn)區(qū)別：

• 作為Module參數(shù)會(huì)被自動(dòng)加入到該Module的參數(shù)列表中；
• 不能被volatile， 默認(rèn)require gradient。

Module

Module為所有神經(jīng)網(wǎng)絡(luò)模塊的父類(lèi)，如開(kāi)始的例子，Net繼承該類(lèi)，____init____中指定網(wǎng)絡(luò)結(jié)構(gòu)中的模塊，并重寫(xiě)forward方法實(shí)現(xiàn)前向傳播得到指定輸入的輸出值，以此進(jìn)行后面loss的計(jì)算和反向傳播。

Optimizer

Optimizer是所有優(yōu)化算法的父類(lèi)（SGD、Adam、...），____init____中傳入網(wǎng)絡(luò)的parameters, 子類(lèi)實(shí)現(xiàn)父類(lèi)step方法，完成對(duì)parameters的更新。

自定義Module

該部分說(shuō)明自定義的Module是如何組織定義在構(gòu)造函數(shù)中的子Module,以及自定義的parameters的保存形式，eg:

1. class Net(nn.Module):
2. def __init__(self):
3. super(Net, self).__init__()
4. self.linear = nn.Linear(4, 2)
6. def forward(self, input):
7. out = F.softmax(self.linear(input))
8. return out

首先看構(gòu)造函數(shù)，Module的構(gòu)造函數(shù)初始化了Module的基本屬性，這里關(guān)注_parameters和_modules，兩個(gè)屬性初始化為OrderedDict()，pytorch重寫(xiě)的有序字典類(lèi)型。_parameters保存網(wǎng)絡(luò)的所有參數(shù)，_modules保存當(dāng)前Module的子Module。
module.py：

1. class Module(object):
3. def __init__(self):
4. self._parameters = OrderedDict()
5. self._modules = OrderedDict()
6. ...

下面來(lái)看自定義Net類(lèi)中self.linear = nn.Linear(4, 2)語(yǔ)句和_modules、_parameters如何產(chǎn)生聯(lián)系，或者self.linear及其參數(shù)如何被添加到_modules、_parameters字典中。答案在Module的____setattr____方法，該P(yáng)ython內(nèi)建方法會(huì)在類(lèi)的屬性被賦值時(shí)調(diào)用。
module.py:

1. def __setattr__(self, name, value):
2. def remove_from(*dicts):
3. for d in dicts:
4. if name in d:
5. del d[name]
7. params = self.__dict__.get('_parameters')
8. if isinstance(value, Parameter): # ----------- <1>
9. if params is None:
10. raise AttributeError(
11. "cannot assign parameters before Module.__init__() call")
12. remove_from(self.__dict__, self._buffers, self._modules)
13. self.register_parameter(name, value)
14. elif params is not None and name in params:
15. if value is not None:
16. raise TypeError("cannot assign '{}' as parameter '{}' "
17. "(torch.nn.Parameter or None expected)"
18. .format(torch.typename(value), name))
19. self.register_parameter(name, value)
20. else:
21. modules = self.__dict__.get('_modules')
22. if isinstance(value, Module):# ----------- <2>
23. if modules is None:
24. raise AttributeError(
25. "cannot assign module before Module.__init__() call")
26. remove_from(self.__dict__, self._parameters, self._buffers)
27. modules[name] = value
28. elif modules is not None and name in modules:
29. if value is not None:
30. raise TypeError("cannot assign '{}' as child module '{}' "
31. "(torch.nn.Module or None expected)"
32. .format(torch.typename(value), name))
33. modules[name] = value
34. ......

調(diào)用self.linear = nn.Linear(4, 2)時(shí)，父類(lèi)____setattr____被調(diào)用，參數(shù)name為“l(fā)inear”， value為nn.Linear(4, 2)，內(nèi)建的Linear類(lèi)同樣是Module的子類(lèi)。所以<2>中的判斷為真，接著modules[name] = value，該linear被加入_modules字典。
同樣自定義Net類(lèi)的參數(shù)即為其子模塊Linear的參數(shù)，下面看Linear的實(shí)現(xiàn)：
linear.py:

1. class Linear(Module):
3. def __init__(self, in_features, out_features, bias=True):
4. super(Linear, self).__init__()
5. self.in_features = in_features
6. self.out_features = out_features
7. self.weight = Parameter(torch.Tensor(out_features, in_features))
8. if bias:
9. self.bias = Parameter(torch.Tensor(out_features))
10. else:
11. self.register_parameter('bias', None)
12. self.reset_parameters()
14. def reset_parameters(self):
15. stdv = 1. / math.sqrt(self.weight.size(1))
16. self.weight.data.uniform_(-stdv, stdv)
17. if self.bias is not None:
18. self.bias.data.uniform_(-stdv, stdv)
20. def forward(self, input):
21. return F.linear(input, self.weight, self.bias)

同樣繼承Module類(lèi)，____init____中參數(shù)為輸入輸出維度，是否需要bias參數(shù)。在self.weight = Parameter(torch.Tensor(out_features, in_features))的初始化時(shí)，同樣會(huì)調(diào)用父類(lèi)Module的____setattr____， name為“weight”，value為Parameter，此時(shí)<1>判斷為真，調(diào)用self.register_parameter(name, value)，該方法中對(duì)參數(shù)進(jìn)行合法性校驗(yàn)后放入self._parameters字典中。

Linear在reset_parameters方法對(duì)權(quán)重進(jìn)行了初始化。

最終可以得出結(jié)論自定義的Module以樹(shù)的形式組織子Module，子Module及其參數(shù)以字典的方式保存。

前向傳播、反向傳播

前向傳播

例子中out = net(features)實(shí)現(xiàn)了網(wǎng)絡(luò)的前向傳播，該語(yǔ)句會(huì)調(diào)用Module類(lèi)的forward方法，該方法被繼承父類(lèi)的子類(lèi)實(shí)現(xiàn)。net(features)使用對(duì)象作為函數(shù)調(diào)用，會(huì)調(diào)用Python內(nèi)建的____call____方法，Module重寫(xiě)了該方法。

module.py:

1. def __call__(self, *input, **kwargs):
2. for hook in self._forward_pre_hooks.values():
3. hook(self, input)
4. result = self.forward(*input, **kwargs)
5. for hook in self._forward_hooks.values():
6. hook_result = hook(self, input, result)
7. if hook_result is not None:
8. raise RuntimeError(
9. "forward hooks should never return any values, but '{}'"
10. "didn't return None".format(hook))
11. if len(self._backward_hooks) > 0:
12. var = result
13. while not isinstance(var, Variable):
14. var = var[0]
15. grad_fn = var.grad_fn
16. if grad_fn is not None:
17. for hook in self._backward_hooks.values():
18. wrapper = functools.partial(hook, self)
19. functools.update_wrapper(wrapper, hook)
20. grad_fn.register_hook(wrapper)
21. return result

____call____方法中調(diào)用result = self.forward(*input, **kwargs)前后會(huì)查看有無(wú)hook函數(shù)需要調(diào)用（預(yù)處理和后處理）。
例子中Net的forward方法中out = F.softmax(self.linear(input))，同樣會(huì)調(diào)用self.linear的forward方法F.linear(input, self.weight, self.bias)進(jìn)行矩陣運(yùn)算（仿射變換）。
functional.py:

1. def linear(input, weight, bias=None):
2. if input.dim() == 2 and bias is not None:
3. # fused op is marginally faster
4. return torch.addmm(bias, input, weight.t())
6. output = input.matmul(weight.t())
7. if bias is not None:
8. output += bias
9. return output

最終經(jīng)過(guò)F.softmax，得到前向輸出結(jié)果。F.softmax和F.linear類(lèi)似前面說(shuō)到的Function（Parameters的表達(dá)式或計(jì)算圖）。

反向傳播

得到前向傳播結(jié)果后，計(jì)算loss = F.cross_entropy(out, target)，接下來(lái)反向傳播求導(dǎo)數(shù)d(loss)/d(weight)和d(loss)/d(bias)：

loss.backward()

backward()方法同樣底層由C拓展，這里暫不深入，調(diào)用該方法后，loss計(jì)算圖中的所有Variable(這里linear的weight和bias)的grad被求出。

Optimizer參數(shù)更新

在計(jì)算出參數(shù)的grad后，需要根據(jù)優(yōu)化算法對(duì)參數(shù)進(jìn)行更新，不同的優(yōu)化算法有不同的更新策略。
optimizer.py:

1. class Optimizer(object):
3. def __init__(self, params, defaults):
4. if isinstance(params, Variable) or torch.is_tensor(params):
5. raise TypeError("params argument given to the optimizer should be "
6. "an iterable of Variables or dicts, but got " +
7. torch.typename(params))
9. self.state = defaultdict(dict)
10. self.param_groups = list(params)
11. ......
13. def zero_grad(self):
14. """Clears the gradients of all optimized :class:`Variable` s."""
15. for group in self.param_groups:
16. for p in group['params']:
17. if p.grad is not None:
18. if p.grad.volatile:
19. p.grad.data.zero_()
20. else:
21. data = p.grad.data
22. p.grad = Variable(data.new().resize_as_(data).zero_())
24. def step(self, closure):
25. """Performs a single optimization step (parameter update).
27. Arguments:
28. closure (callable): A closure that reevaluates the model and
29. returns the loss. Optional for most optimizers.
30. """
31. raise NotImplementedError

Optimizer在init中將傳入的params保存到self.param_groups，另外兩個(gè)重要的方法zero_grad負(fù)責(zé)將參數(shù)的grad置零方便下次計(jì)算，step負(fù)責(zé)參數(shù)的更新，由子類(lèi)實(shí)現(xiàn)。
以列子中的sgd = torch.optim.SGD(net.parameters(), lr=0.001)為例，其中net.parameters()返回Net參數(shù)的迭代器，為待優(yōu)化參數(shù)；lr指定學(xué)習(xí)率。
SGD.py:

1. class SGD(Optimizer):
3. def __init__(self, params, lr=required, momentum=0, dampening=0,
4. weight_decay=0, nesterov=False):
5. defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
6. weight_decay=weight_decay, nesterov=nesterov)
7. if nesterov and (momentum <= 0 or dampening != 0):
8. raise ValueError("Nesterov momentum requires a momentum and zero dampening")
9. super(SGD, self).__init__(params, defaults)
11. def __setstate__(self, state):
12. super(SGD, self).__setstate__(state)
13. for group in self.param_groups:
14. group.setdefault('nesterov', False)
16. def step(self, closure=None):
17. """Performs a single optimization step.
19. Arguments:
20. closure (callable, optional): A closure that reevaluates the model
21. and returns the loss.
22. """
23. loss = None
24. if closure is not None:
25. loss = closure()
27. for group in self.param_groups:
28. weight_decay = group['weight_decay']
29. momentum = group['momentum']
30. dampening = group['dampening']
31. nesterov = group['nesterov']
33. for p in group['params']:
34. if p.grad is None:
35. continue
36. d_p = p.grad.data
37. if weight_decay != 0:
38. d_p.add_(weight_decay, p.data)
39. if momentum != 0:
40. param_state = self.state[p]
41. if 'momentum_buffer' not in param_state:
42. buf = param_state['momentum_buffer'] = d_p.clone()
43. else:
44. buf = param_state['momentum_buffer']
45. buf.mul_(momentum).add_(1 - dampening, d_p)
46. if nesterov:
47. d_p = d_p.add(momentum, buf)
48. else:
49. d_p = buf
51. p.data.add_(-group['lr'], d_p)
53. return loss

SGD的step方法中，判斷是否使用權(quán)重衰減和動(dòng)量更新，如果不使用，直接更新權(quán)重param := param - lr * d(param)。例子中調(diào)用sgd.step()后完成一次epoch。這里由于傳遞到Optimizer的參數(shù)集是可更改（mutable）的，step中對(duì)參數(shù)的更新同樣是Net中參數(shù)的更新。

小結(jié)

到此，根據(jù)一個(gè)簡(jiǎn)單例子閱讀了pytorch中Python實(shí)現(xiàn)的部分源碼，沒(méi)有深入到底層Tensor、autograd等部分的C拓展實(shí)現(xiàn)，后面再繼續(xù)讀一讀C拓展部分的代碼。

轉(zhuǎn)自鏈接：https://www.jianshu.com/p/f5eb8c2e671c

總結(jié)

以上是生活随笔為你收集整理的pytorch源码解析：Python层 pytorchmodule源码的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。