@@ -6,18 +6,18 @@ If you haven't, then you might prefer the [Fitting a Straight Line](overview.md)
6
6
7
7
``` julia
8
8
# With Julia 1.7+, this will prompt if neccessary to install everything, including CUDA:
9
- using Flux, Statistics
9
+ using Flux, Zygote, Statistics
10
10
11
11
# Generate some data for the XOR problem: vectors of length 2, as columns of a matrix:
12
12
noisy = rand (Float32, 2 , 1000 ) # 2×1000 Matrix{Float32}
13
13
truth = [xor (col[1 ]> 0.5 , col[2 ]> 0.5 ) for col in eachcol (noisy)] # 1000-element Vector{Bool}
14
14
15
15
# Define our model, a multi-layer perceptron with one hidden layer of size 3:
16
16
model = Chain (
17
- Dense (2 => 3 , tanh), # activation function inside...
17
+ Dense (2 => 3 , tanh), # activation function inside layer
18
18
BatchNorm (3 ),
19
19
Dense (3 => 2 ),
20
- softmax) # ... but softmax outside a layer.
20
+ softmax)
21
21
22
22
# The model encapsulates parameters, randomly initialised. Its initial output is:
23
23
out1 = model (noisy) # 2×1000 Matrix{Float32}
@@ -34,15 +34,13 @@ opt = Flux.Adam(0.01) # will store optimiser momentum, etc.
34
34
for epoch in 1 : 1_000
35
35
losses = []
36
36
for (x, y) in loader
37
- loss, grad = Flux . withgradient (pars) do
37
+ loss, grad = Zygote . withgradient (pars) do
38
38
# Evaluate model and loss inside gradient context:
39
39
y_hat = model (x)
40
- Flux. crossentropy (y_hat, y) # could use just sum(abs2, y_hat .- y)
40
+ Flux. crossentropy (y_hat, y)
41
41
end
42
- # Use the gradient to update the model's parameters (and momentum):
43
42
Flux. update! (opt, pars, grad)
44
- # Logging code, outside gradient context:
45
- push! (losses, loss)
43
+ push! (losses, loss) # logging, outside gradient context
46
44
end
47
45
if isinteger (log2 (epoch))
48
46
println (" after epoch $epoch , loss is " , mean (losses))
0 commit comments