forked from PyTorchKorea/tutorials-kr
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathforward_ad_usage.py
More file actions
227 lines (180 loc) Β· 11.2 KB
/
forward_ad_usage.py
File metadata and controls
227 lines (180 loc) Β· 11.2 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
# -*- coding: utf-8 -*-
"""
μμ ν λͺ¨λ μλ λ―ΈλΆ(Beta)
=============================================
**λ²μ**: `κΉκ²½λ―Ό <https://github.com/BcKmini>`_
μ΄ νν 리μΌμ μμ ν λͺ¨λ μλ λ―ΈλΆ(Forward-mode Automatic Differentiation)μ μ¬μ©νμ¬ λ°©ν₯μ± λν¨μ(directional derivative) λλ μΌμ½λΉμ-λ²‘ν° κ³±(Jacobian-vector product)μ κ³μ°νλ λ°©λ²μ 보μ¬μ€λλ€.
μλ νν 리μΌμ 1.11 μ΄μ λ²μ (λλ λμ΄ν리 λΉλ)μμλ§ μ¬μ©ν μ μλ μΌλΆ APIλ₯Ό μ¬μ©ν©λλ€.
λν, μμ ν λͺ¨λ μλ λ―ΈλΆμ νμ¬ λ² ν λ²μ μ
λλ€. λ°λΌμ APIκ° λ³κ²½λ μ μμΌλ©°, μμ§ μΌλΆ μ°μ°μλ μ§μλμ§ μμ μ μμ΅λλ€.
κΈ°λ³Έ μ¬μ©λ²
--------------------------------------------------------------------
μμ ν λͺ¨λ μλ λ―ΈλΆ(Reverse-mode Automatic Differentiation)κ³Ό λ¬λ¦¬, μμ ν λͺ¨λ μλ λ―ΈλΆμ μμ ν(forward pass)λ₯Ό μ§ννλ©° κΈ°μΈκΈ°(gradient)λ₯Ό μ¦μ(κ³μ°μ λ―Έλ£¨μ§ μκ³ ) κ³μ°ν©λλ€. μμ ν λͺ¨λ μλ λ―ΈλΆμΌλ‘ λ°©ν₯μ± λν¨μλ₯Ό κ³μ°νλ €λ©΄, λ¨Όμ μ
λ ₯μ λ°©ν₯μ± λν¨μμ λ°©ν₯μ λνλ΄λ λ€λ₯Έ ν
μ(μΌμ½λΉμ-λ²‘ν° κ³±μ `v`μ ν΄λΉ)μ μ°κ²°ν λ€ μ΄μ κ³Ό κ°μ΄ μμ νλ₯Ό μννλ©΄ λ©λλ€. 'primal'μ΄λΌκ³ λΆλ₯΄λ μ
λ ₯μ΄ 'tangent'λΌκ³ λΆλ₯΄λ 'λ°©ν₯' ν
μμ μ°κ²°λ λ, κ²°κ³Όλ‘ λμ€λ μλ‘μ΄ ν
μ κ°μ²΄λ μ΄μ€μ(dual numbers) [0] μμ κ΄λ ¨μ± λλ¬Έμ 'μ΄μ€ ν
μ(dual tensor)'λΌκ³ λΆλ¦½λλ€.
μμ νκ° μνλ λ, μ
λ ₯ ν
μ μ€ νλλΌλ μ΄μ€ ν
μμ΄λ©΄ ν¨μμ 'λ―Όκ°λ(sensitivity)'λ₯Ό μ ννκΈ° μν΄ μΆκ°μ μΈ μ°μ°μ΄ μνλ©λλ€.
"""
import torch
import torch.autograd.forward_ad as fwAD
primal = torch.randn(10, 10)
tangent = torch.randn(10, 10)
def fn(x, y):
return x ** 2 + y ** 2
# λͺ¨λ μμ ν μλ λ―ΈλΆ μ°μ°μ ``dual_level`` 컨ν
μ€νΈ μμμ μνν΄μΌ ν©λλ€.
# μ΄ μ»¨ν
μ€νΈμμ μμ±λ λͺ¨λ μ΄μ€ ν
μμ νμ νΈλ 컨ν
μ€νΈλ₯Ό λ²μ΄λ λ μλ©Έλ©λλ€.
# μ΄λ ν΄λΉ μ°μ°μ μΆλ ₯μ΄λ μ€κ° κ²°κ³Όκ° ν₯ν λ€λ₯Έ μμ ν μλ λ―ΈλΆ μ°μ°μ μ¬μ¬μ©λ λ,
# νμ¬ μ°μ°μ μν νμ νΈκ° λμ€ μ°μ°μ νμ νΈμ νΌλλλ κ²μ λ°©μ§νκΈ° μν¨μ
λλ€.
with fwAD.dual_level():
# μ΄μ€ ν
μλ₯Ό λ§λ€λ €λ©΄ 'primal' ν
μλ₯Ό κ°μ ν¬κΈ°μ λ€λ₯Έ ν
μ,
# μ¦ 'νμ νΈ(tangent)'μ μ°κ²°ν©λλ€.
# λ§μ½ νμ νΈμ λ μ΄μμμ΄ primalμ λ μ΄μμκ³Ό λ€λ₯΄λ©΄,
# νμ νΈμ κ°μ primalκ³Ό λμΌν λ©νλ°μ΄ν°λ₯Ό κ°λ μ ν
μμ 볡μ¬λ©λλ€.
# κ·Έλ μ§ μμΌλ©΄ νμ νΈ μμ²΄κ° κ·Έλλ‘ μ¬μ©λ©λλ€.
#
# ``make_dual`` λ‘ μμ±λ μ΄μ€ ν
μλ primal ν
μμ **λ·°(λ°μ΄ν°λ₯Ό 곡μ νλ μ°Έμ‘°)** λΌλ μ λ
# μ€μν©λλ€.
dual_input = fwAD.make_dual(primal, tangent)
assert fwAD.unpack_dual(dual_input).tangent is tangent
# νμ νΈκ° 볡μ¬λλ κ²½μ°λ₯Ό 보μ¬μ£ΌκΈ° μν΄,
# primalκ³Ό λ€λ₯Έ λ μ΄μμμ κ°μ§ νμ νΈλ₯Ό μ λ¬ν©λλ€.
dual_input_alt = fwAD.make_dual(primal, tangent.T)
assert fwAD.unpack_dual(dual_input_alt).tangent is not tangent
# νμ νΈκ° μ°κ²°λμ§ μμ ν
μλ μλμΌλ‘
# κ°μ shapeμ κ°μ§λ©° 0μΌλ‘ μ±μμ§ νμ νΈλ₯Ό κ°μ§ κ²μΌλ‘ κ°μ£Όλ©λλ€.
plain_tensor = torch.randn(10, 10)
dual_output = fn(dual_input, plain_tensor)
# μ΄μ€ ν
μλ₯Ό νλ©΄(unpack) ``primal`` κ³Ό ``tangent`` λ₯Ό
# μμ±μΌλ‘ κ°λ ``namedtuple`` μ΄ λ°νλ©λλ€.
jvp = fwAD.unpack_dual(dual_output).tangent
assert fwAD.unpack_dual(dual_output).tangent is None
######################################################################
# λͺ¨λκ³Ό ν¨κ» μ¬μ©νκΈ°
# --------------------------------------------------------------------
# ``nn.Module`` μ μμ ν μλ λ―ΈλΆκ³Ό ν¨κ» μ¬μ©νλ €λ©΄, μμ νλ₯Ό μννκΈ° μ μ
# λͺ¨λΈμ λ§€κ°λ³μ(parameter)λ₯Ό μ΄μ€ ν
μλ‘ κ΅μ²΄ν΄μΌ ν©λλ€. νμ¬ μ΄μ€ ν
μλ‘ λ
# `nn.Parameter` λ μμ±ν μ μμ΅λλ€. μ΄μ λν ν΄κ²° λ°©λ²μΌλ‘,
# μ΄μ€ ν
μλ₯Ό λͺ¨λμ λ§€κ°λ³μκ° μλ μΌλ° μμ±μΌλ‘ λ±λ‘ν΄μΌ ν©λλ€.
import torch.nn as nn
model = nn.Linear(5, 5)
input = torch.randn(16, 5)
params = {name: p for name, p in model.named_parameters()}
tangents = {name: torch.rand_like(p) for name, p in params.items()}
with fwAD.dual_level():
for name, p in params.items():
delattr(model, name)
setattr(model, name, fwAD.make_dual(p, tangents[name]))
out = model(input)
jvp = fwAD.unpack_dual(out).tangent
######################################################################
# ν¨μν λͺ¨λ API μ¬μ©νκΈ° (Beta)
# --------------------------------------------------------------------
# ``nn.Module`` μ μμ ν μλ λ―ΈλΆκ³Ό ν¨κ» μ¬μ©νλ λ λ€λ₯Έ λ°©λ²μ
# ν¨μν λͺ¨λ APIλ₯Ό νμ©νλ κ²μ
λλ€.
from torch.func import functional_call
# functional_callμ λͺ¨λΈμ λ§€κ°λ³μκ° λ±λ‘λμ΄ μμ΄μΌ νλ―λ‘
# μλ‘μ΄ λͺ¨λμ΄ νμν©λλ€.
model = nn.Linear(5, 5)
dual_params = {}
with fwAD.dual_level():
for name, p in params.items():
# μ μΉμ
κ³Ό λμΌν ``tangents`` λ₯Ό μ¬μ©ν©λλ€.
dual_params[name] = fwAD.make_dual(p, tangents[name])
out = functional_call(model, dual_params, input)
jvp2 = fwAD.unpack_dual(out).tangent
# κ²°κ³Ό νμΈ
assert torch.allclose(jvp, jvp2)
######################################################################
# μ¬μ©μ μ μ autograd ν¨μ
# --------------------------------------------------------------------
# μ¬μ©μ μ μ ν¨μ λν μμ ν λͺ¨λ μλ λ―ΈλΆμ μ§μν©λλ€. μμ ν λͺ¨λ μλ λ―ΈλΆμ
# μ§μνλ μ¬μ©μ μ μ ν¨μλ₯Ό λ§λ€λ €λ©΄ ``jvp()`` μ μ λ©μλλ₯Ό
# λ±λ‘ν΄μΌ ν©λλ€. μ¬μ©μ μ μ ν¨μκ° μμ νμ μμ ν μλ λ―ΈλΆμ λͺ¨λ μ§μνλ κ²λ
# κ°λ₯νμ§λ§ νμλ μλλλ€. λ μμΈν μ 보λ
# `λ¬Έμ <https://pytorch.org/docs/master/notes/extending.html#forward-mode-ad>`_
# λ₯Ό μ°Έκ³ νμΈμ.
class Fn(torch.autograd.Function):
@staticmethod
def forward(ctx, foo):
result = torch.exp(foo)
# ``ctx`` μ μ μ₯λ ν
μλ μ΄νμ μμ ν κΈ°μΈκΈ°
# κ³μ°μ μ¬μ©ν μ μμ΅λλ€.
ctx.result = result
return result
@staticmethod
def jvp(ctx, gI):
gO = gI * ctx.result
# ``ctx`` μ μ μ₯λ ν
μκ° μμ νμ μ¬μ©λμ§ μμ κ²½μ°,
# ``del`` μ μ¬μ©νμ¬ μλμΌλ‘ λ©λͺ¨λ¦¬μμ ν΄μ ν μ μμ΅λλ€.
del ctx.result
return gO
fn = Fn.apply
primal = torch.randn(10, 10, dtype=torch.double, requires_grad=True)
tangent = torch.randn(10, 10)
with fwAD.dual_level():
dual_input = fwAD.make_dual(primal, tangent)
dual_output = fn(dual_input)
jvp = fwAD.unpack_dual(dual_output).tangent
# μ¬μ©μ μ μ autograd ν¨μκ° κΈ°μΈκΈ°λ₯Ό μ¬λ°λ₯΄κ² κ³μ°νλμ§ νμΈνλ €λ©΄
# ``autograd.gradcheck`` λ₯Ό μ¬μ©νλ κ²μ΄ μ€μν©λλ€. κΈ°λ³Έμ μΌλ‘
# ``gradcheck`` λ μμ ν λͺ¨λ(reverse-mode) μλ λ―ΈλΆ κΈ°μΈκΈ°λ§ νμΈν©λλ€.
# ``check_forward_ad=True`` λ₯Ό μ§μ νμ¬ μμ ν κΈ°μΈκΈ°λ νμΈνλλ‘ ν μ μμ΅λλ€.
# λ§μ½ ν¨μμ λν μμ νλ₯Ό ꡬννμ§ μμλ€λ©΄, ``check_backward_ad=False``,
# ``check_undefined_grad=False``, ``check_batched_grad=False`` λ₯Ό μ§μ νμ¬
# ``gradcheck`` κ° μμ ν λͺ¨λ μλ λ―ΈλΆμ΄ νμν ν
μ€νΈλ₯Ό 건λλ°λλ‘ ν μ μμ΅λλ€.
torch.autograd.gradcheck(Fn.apply, (primal,), check_forward_ad=True,
check_backward_ad=False, check_undefined_grad=False,
check_batched_grad=False)
######################################################################
# ν¨μν API (Beta)
# --------------------------------------------------------------------
# Functorchλ μΌμ½λΉμ-λ²‘ν° κ³±μ κ³μ°νκΈ° μν κ³ μμ€ ν¨μν APIλ
# μ 곡νλ©°, μ¬μ© μ¬λ‘μ λ°λΌ λ κ°λ¨νκ² μ¬μ©ν μ μμ΅λλ€.
#
# ν¨μν APIμ μ₯μ μ μ μμ€μ μ΄μ€ ν
μ APIλ₯Ό μ΄ν΄νκ±°λ μ¬μ©ν
# νμκ° μμΌλ©°, λ€λ₯Έ `functorch λ³ν(vmap λ±)κ³Ό κ²°ν© <https://pytorch.org/functorch/stable/notebooks/jacobians_hessians.html>`_
# ν μ μλ€λ κ²μ
λλ€. λ¨μ μ μΈλ°ν μ μ΄κ° μ΄λ ΅λ€λ μ μ
λλ€.
#
# μ΄ νν 리μΌμ λλ¨Έμ§ λΆλΆμ μ€ννλ €λ©΄ functorch
# (https://github.com/pytorch/functorch) κ° νμν©λλ€.
# μ€μΉ λ°©λ²μ ν΄λΉ λ§ν¬μμ νμΈν΄μ£ΌμΈμ.
import functorch as ft
primal0 = torch.randn(10, 10)
tangent0 = torch.randn(10, 10)
primal1 = torch.randn(10, 10)
tangent1 = torch.randn(10, 10)
def fn(x, y):
return x ** 2 + y ** 2
# μ ν¨μμ JVPλ₯Ό κ³μ°νλ κΈ°λ³Έ μμ μ
λλ€.
# ``jvp(func, primals, tangents)`` λ ``func(*primals)`` μ κ²°κ³Όμ κ³μ°λ
# μΌμ½λΉμ-λ²‘ν° κ³±(JVP)μ ν¨κ» λ°νν©λλ€. κ° primalμ κ°μ shapeμ νμ νΈμ
# μ°κ²°λμ΄μΌ ν©λλ€.
primal_out, tangent_out = ft.jvp(fn, (primal0, primal1), (tangent0, tangent1))
# ``functorch.jvp`` λ λͺ¨λ primalμ΄ νμ νΈμ μ°κ²°λ κ²μ μꡬν©λλ€.
# λ§μ½ ``fn`` μ νΉμ μ
λ ₯μλ§ νμ νΈλ₯Ό μ°κ²°νκ³ μΆλ€λ©΄,
# νμ νΈκ° μλ μ
λ ₯μ λ°λ μλ‘μ΄ ν¨μλ₯Ό λ§λ€μ΄μΌ ν©λλ€.
primal = torch.randn(10, 10)
tangent = torch.randn(10, 10)
y = torch.randn(10, 10)
import functools
new_fn = functools.partial(fn, y=y)
primal_out, tangent_out = ft.jvp(new_fn, (primal,), (tangent,))
######################################################################
# ν¨μν APIλ₯Ό λͺ¨λκ³Ό ν¨κ» μ¬μ©νκΈ°
# --------------------------------------------------------------------
# ``nn.Module`` κ³Ό ``functorch.jvp`` λ₯Ό ν¨κ» μ¬μ©νμ¬ λͺ¨λΈ λ§€κ°λ³μμ λν
# μΌμ½λΉμ-λ²‘ν° κ³±μ κ³μ°νλ €λ©΄, ``nn.Module`` μ λͺ¨λΈ λ§€κ°λ³μμ
# λͺ¨λμ μ
λ ₯μ λͺ¨λ μΈμλ‘ λ°λ ν¨μλ‘ μ¬κ΅¬μ±ν΄μΌ ν©λλ€.
model = nn.Linear(5, 5)
input = torch.randn(16, 5)
tangents = tuple([torch.rand_like(p) for p in model.parameters()])
# ``ft.make_functional_with_buffers`` λ μ£Όμ΄μ§ ``torch.nn.Module`` μμ
# μν(``params`` μ λ²νΌ)λ₯Ό μΆμΆνκ³ , ν¨μμ²λΌ νΈμΆν μ μλ
# ν¨μν λ²μ μ λͺ¨λΈμ λ°νν©λλ€.
# μ¦, λ°νλ ``func`` λ ``func(params, buffers, input)`` μ²λΌ νΈμΆν μ μμ΅λλ€.
# ``ft.make_functional_with_buffers`` λ μ΄μ μ 보μλ ``nn.Module`` μ μν μλ APIμ
# μ μ¬νλ©°, μ΄ λμ ν΅ν©νλ μμ
μ΄ μ§ν μ€μ
λλ€.
func, params, buffers = ft.make_functional_with_buffers(model)
# ``jvp`` λ λͺ¨λ μ
λ ₯μ΄ νμ νΈμ μ°κ²°λ κ²μ μꡬνλ―λ‘,
# λ§€κ°λ³μλ₯Ό λ°μμ λ μΆλ ₯μ μμ±νλ μλ‘μ΄ ν¨μλ₯Ό λ§λ€μ΄μΌ ν©λλ€.
def func_params_only(params):
return func(params, buffers, input)
model_output, jvp_out = ft.jvp(func_params_only, (params,), (tangents,))
######################################################################
# [0] https://en.wikipedia.org/wiki/Dual_number