Skip to content

Commit fd54428

Browse files
committed
Add custom model example to docs.
1 parent f66be89 commit fd54428

File tree

1 file changed

+29
-1
lines changed

1 file changed

+29
-1
lines changed

docs/src/models/advanced.md

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,34 @@
22

33
Here we will try and describe usage of some more advanced features that Flux provides to give more control over model building.
44

5+
## Custom Model Example
6+
7+
Here is a basic example of a custom model. It simply adds the input to the result from the neural network.
8+
9+
```julia
10+
struct CustomModel
11+
chain::Chain
12+
end
13+
14+
function (m::CustomModel)(x)
15+
return m.chain(x) + x
16+
17+
# You can put arbitrary code here, but note that everything here will be differentiated.
18+
# Zygote does not allow some operations, like mutating arrays.
19+
end
20+
21+
# Call @functor to allow for training. Described below in more detail.
22+
Flux.@functor CustomModel
23+
```
24+
25+
You can then use the model like:
26+
27+
```julia
28+
chain = Chain(Dense(10, 10))
29+
model = CustomModel(chain)
30+
model(rand(10))
31+
```
32+
533
## Customising Parameter Collection for a Model
634

735
Taking reference from our example `Affine` layer from the [basics](basics.md#Building-Layers-1).
@@ -68,7 +96,7 @@ by simply deleting it from `ps`:
6896

6997
```julia
7098
ps = params(m)
71-
delete!(ps, m[2].bias)
99+
delete!(ps, m[2].bias)
72100
```
73101

74102
## Custom multiple input or output layer

0 commit comments

Comments
 (0)