@@ -4,7 +4,7 @@ We will train a simple Deep Equilibrium Model on MNIST. First we load a few pack
44
55``` @example basic_mnist_deq
66using DeepEquilibriumNetworks, SciMLSensitivity, Lux, NonlinearSolve, OrdinaryDiffEq,
7- Statistics, Random, Optimisers, LuxCUDA, Zygote, LinearSolve, LoggingExtras
7+ Statistics, Random, Optimisers, LuxCUDA, Zygote, LinearSolve, Dates, Printf
88using MLDatasets: MNIST
99using MLDataUtils: LabelEnc, convertlabel, stratifiedobs, batchview
1010
@@ -20,18 +20,6 @@ const cdev = cpu_device()
2020const gdev = gpu_device()
2121```
2222
23- SciMLBase introduced a warning instead of depwarn which pollutes the output. We can suppress
24- it with the following logger
25-
26- ``` @example basic_mnist_deq
27- function remove_syms_warning(log_args)
28- return log_args.message !=
29- "The use of keyword arguments `syms`, `paramsyms` and `indepsym` for `SciMLFunction`s is deprecated. Pass `sys = SymbolCache(syms, paramsyms, indepsym)` instead."
30- end
31-
32- filtered_logger = ActiveFilteredLogger(remove_syms_warning, global_logger())
33- ```
34-
3523We can now construct our dataloader.
3624
3725``` @example basic_mnist_deq
@@ -66,8 +54,7 @@ function construct_model(solver; model_type::Symbol=:deq)
6654
6755 # The input layer of the DEQ
6856 deq_model = Chain(
69- Parallel(+,
70- Conv((3, 3), 64 => 64, tanh; stride=1, pad=SamePad()),
57+ Parallel(+, Conv((3, 3), 64 => 64, tanh; stride=1, pad=SamePad()),
7158 Conv((3, 3), 64 => 64, tanh; stride=1, pad=SamePad())),
7259 Conv((3, 3), 64 => 64, tanh; stride=1, pad=SamePad()))
7360
@@ -79,11 +66,11 @@ function construct_model(solver; model_type::Symbol=:deq)
7966 init = missing
8067 end
8168
82- deq = DeepEquilibriumNetwork(deq_model, solver; init, verbose=false,
83- linsolve_kwargs=(; maxiters=10))
69+ deq = DeepEquilibriumNetwork(
70+ deq_model, solver; init, verbose=false, linsolve_kwargs=(; maxiters=10))
8471
85- classifier = Chain(GroupNorm(64, 64, relu), GlobalMeanPool(), FlattenLayer(),
86- Dense(64, 10))
72+ classifier = Chain(
73+ GroupNorm(64, 64, relu), GlobalMeanPool(), FlattenLayer(), Dense(64, 10))
8774
8875 model = Chain(; down, deq, classifier)
8976
@@ -95,12 +82,12 @@ function construct_model(solver; model_type::Symbol=:deq)
9582 x = randn(rng, Float32, 28, 28, 1, 128)
9683 y = onehot(rand(Random.default_rng(), 0:9, 128)) |> gdev
9784
98- model_ = Lux.Experimental. StatefulLuxLayer(model, ps, st)
99- @info " warming up forward pass"
85+ model_ = StatefulLuxLayer(model, ps, st)
86+ @printf "[%s] warming up forward pass\n" string(now())
10087 logitcrossentropy(model_, x, ps, y)
101- @info " warming up backward pass"
88+ @printf "[%s] warming up backward pass\n" string(now())
10289 Zygote.gradient(logitcrossentropy, model_, x, ps, y)
103- @info " warmup complete"
90+ @printf "[%s] warmup complete\n" string(now())
10491
10592 return model, ps, st
10693end
@@ -122,7 +109,7 @@ classify(x) = argmax.(eachcol(x))
122109function accuracy(model, data, ps, st)
123110 total_correct, total = 0, 0
124111 st = Lux.testmode(st)
125- model = Lux.Experimental. StatefulLuxLayer(model, ps, st)
112+ model = StatefulLuxLayer(model, ps, st)
126113 for (x, y) in data
127114 target_class = classify(cdev(y))
128115 predicted_class = classify(cdev(model(x)))
@@ -132,51 +119,48 @@ function accuracy(model, data, ps, st)
132119 return total_correct / total
133120end
134121
135- function train_model(solver, model_type; data_train=zip(x_train, y_train),
136- data_test=zip(x_test, y_test))
122+ function train_model(
123+ solver, model_type; data_train=zip(x_train, y_train), data_test=zip(x_test, y_test))
137124 model, ps, st = construct_model(solver; model_type)
138- model_st = Lux.Experimental. StatefulLuxLayer(model, nothing, st)
125+ model_st = StatefulLuxLayer(model, nothing, st)
139126
140- @info " Training Model: $(model_type) with Solver: $( nameof(typeof(solver)))"
127+ @printf "[%s] Training Model: %s with Solver: %s\n" string(now()) model_type nameof(typeof(solver))
141128
142129 opt_st = Optimisers.setup(Adam(0.001), ps)
143130
144131 acc = accuracy(model, data_test, ps, st) * 100
145- @info " Starting Accuracy: $(acc)"
132+ @printf "[%s] Starting Accuracy: %.5f%%\n" string(now()) acc
146133
147- @info " Pretrain with unrolling to a depth of 5"
134+ @printf "[%s] Pretrain with unrolling to a depth of 5\n" string(now())
148135 st = Lux.update_state(st, :fixed_depth, Val(5))
149- model_st = Lux.Experimental. StatefulLuxLayer(model, ps, st)
136+ model_st = StatefulLuxLayer(model, ps, st)
150137
151138 for (i, (x, y)) in enumerate(data_train)
152139 res = Zygote.withgradient(logitcrossentropy, model_st, x, ps, y)
153140 Optimisers.update!(opt_st, ps, res.grad[3])
154- if i % 50 == 1
155- @info "Pretraining Batch: [$(i)/$(length(data_train))] Loss: $(res.val)"
156- end
141+ i % 50 == 1 &&
142+ @printf "[%s] Pretraining Batch: [%4d/%4d] Loss: %.5f\n" string(now()) i length(data_train) res.val
157143 end
158144
159145 acc = accuracy(model, data_test, ps, model_st.st) * 100
160- @info " Pretraining complete. Accuracy: $(acc)"
146+ @printf "[%s] Pretraining complete. Accuracy: %.5f%%\n" string(now()) acc
161147
162148 st = Lux.update_state(st, :fixed_depth, Val(0))
163- model_st = Lux.Experimental. StatefulLuxLayer(model, ps, st)
149+ model_st = StatefulLuxLayer(model, ps, st)
164150
165151 for epoch in 1:3
166152 for (i, (x, y)) in enumerate(data_train)
167153 res = Zygote.withgradient(logitcrossentropy, model_st, x, ps, y)
168154 Optimisers.update!(opt_st, ps, res.grad[3])
169- if i % 50 == 1
170- @info "Epoch: [$(epoch)/3] Batch: [$(i)/$(length(data_train))] Loss: $(res.val)"
171- end
155+ i % 50 == 1 &&
156+ @printf "[%s] Epoch: [%d/%d] Batch: [%4d/%4d] Loss: %.5f\n" string(now()) epoch 3 i length(data_train) res.val
172157 end
173158
174159 acc = accuracy(model, data_test, ps, model_st.st) * 100
175- @info " Epoch: [$(epoch)/3 ] Accuracy: $(acc)"
160+ @printf "[%s] Epoch: [%d/%d ] Accuracy: %.5f%%\n" string(now()) epoch 3 acc
176161 end
177162
178- @info "Training complete."
179- println()
163+ @printf "[%s] Training complete.\n" string(now())
180164
181165 return model, ps, st
182166end
@@ -188,19 +172,15 @@ and end up using solvers like `Broyden`, but we can simply slap in any of the fa
188172from NonlinearSolve.jl. Here we will use Newton-Krylov Method:
189173
190174``` @example basic_mnist_deq
191- with_logger(filtered_logger) do
192- train_model(NewtonRaphson(; linsolve=KrylovJL_GMRES()), :regdeq)
193- end
175+ train_model(NewtonRaphson(; linsolve=KrylovJL_GMRES()), :regdeq);
194176nothing # hide
195177```
196178
197179We can also train a continuous DEQ by passing in an ODE solver. Here we will use ` VCAB3() `
198180which tend to be quite fast for continuous Neural Network problems.
199181
200182``` @example basic_mnist_deq
201- with_logger(filtered_logger) do
202- train_model(VCAB3(), :deq)
203- end
183+ train_model(VCAB3(), :deq);
204184nothing # hide
205185```
206186
0 commit comments