FluxTraining
"""
epoch!(learner, phase[, dataiter])
Train `learner` for one epoch on `dataiter`. Iterates through
`dataiter` and [`step!`](#)s for each batch/item.
If no data iterator is passed in, use `learner.data[phasedataiter(phase)]`.
## Extending
The default implementation iterates over every batch in `dataiter`
and calls [`step!`](#) for each. This behavior can be overloaded
by implementing `epoch!(learner, ::MyPhase, dataiter)`.
If you're implementing a custom `epoch!` method, it is recommended
you make use of [`runepoch`](#) to get begin and end events as well
as proper handling of [`CancelEpochException`](#)s.
See the default implementation for reference.
"""
function
epoch!
(
learner
,
phase
::
Phase
,
dataiter
=
learner
.
data
[
phasedataiter
(
phase
)
]
)
runepoch
(
learner
,
phase
)
do
_
for
batch
in
dataiter
step!
(
learner
,
phase
,
batch
)
end
end
end
"""
step!(learner, phase::Phase, batch)
Run one step of training for `learner` on batch.
Behavior is customized through `phase`.
## Extending
This is a required method for custom [`Phase`](#)s to implement.
To implement `step!`, it is recommended you make use of [`runstep`](#)
to get begin and end events as well as proper handling of
[`CancelStepException`](#)s.
See the implementations of [`TrainingPhase`](#) and [`ValidationPhase`](#)
for reference.
"""
function
step!
end
function
step!
(
learner
,
phase
::
TrainingPhase
,
batch
)
xs
,
ys
=
batch
runstep
(
learner
,
phase
,
(
;
xs
=
xs
,
ys
=
ys
)
)
do
handle
,
state
state
.
grads
=
_gradient
(
learner
.
optimizer
,
learner
.
model
,
learner
.
params
)
do
model
state
.
ŷs
=
model
(
state
.
xs
)
handle
(
LossBegin
(
)
)
state
.
loss
=
learner
.
lossfn
(
state
.
ŷs
,
state
.
ys
)
handle
(
BackwardBegin
(
)
)
return
state
.
loss
end
handle
(
BackwardEnd
(
)
)
learner
.
params
,
learner
.
model
=
_update!
(
learner
.
optimizer
,
learner
.
params
,
learner
.
model
,
state
.
grads
)
end
endHandle both old Flux.jl and new Optimisers.jl optimisers
_gradient
(
f
,
_
,
m
,
_
)
=
gradient
(
f
,
m
)
[
1
]
_gradient
(
f
,
::
Flux
.
Optimise
.
AbstractOptimiser
,
m
,
ps
::
Params
)
=
gradient
(
(
)
->
f
(
m
)
,
ps
)
function
_update!
(
optimizer
::
Flux
.
Optimise
.
AbstractOptimiser
,
params
,
model
,
grads
)
update!
(
optimizer
,
params
,
grads
)
return
params
,
model
end
function
_update!
(
_
,
st
,
model
,
grads
)
st
,
model
=
Optimisers
.
update!
(
st
,
model
,
grads
)
return
st
,
model
end
function
step!
(
learner
,
phase
::
ValidationPhase
,
batch
)
xs
,
ys
=
batch
runstep
(
learner
,
phase
,
(
;
xs
=
xs
,
ys
=
ys
)
)
do
_
,
state
state
.
ŷs
=
learner
.
model
(
state
.
xs
)
state
.
loss
=
learner
.
lossfn
(
state
.
ŷs
,
state
.
ys
)
end
end
"""
runepoch(epochfn, learner, phase)
Run `epochfn` inside the context of an epoch. Calls `epochfn(handle)`
where `handle(e)` can be called to dispatch events.
Takes care of dispatching [`EpochBegin`](#) and [`EpochEnd`](#)
events as well as handling [`CancelEpochException`](#)s.
"""
function
runepoch
(
epochfn
,
learner
,
phase
::
Phase
)
handlefn
(
e
)
=
handle
(
learner
.
callbacks
.
runner
,
e
,
phase
,
learner
)
try
handlefn
(
EpochBegin
(
)
)
epochfn
(
handlefn
)
handlefn
(
EpochEnd
(
)
)
catch
e
if
e
isa
CancelEpochException
@
debug
"
Epoch skipped
"
error
=
e
handlefn
(
EpochEnd
(
)
)
else
rethrow
(
)
end
end
end
"""
runstep(stepfn, learner, phase) -> state
Run `stepfn` inside the context of a step. Calls `stepfn(handle, state)`
where `handle(e)` can be called to dispatch events and `state` is a [`PropDict`](#)
which step data, gradients and losses can be written to. Return `state`.
Takes care of dispatching [`StepBegin`](#) and [`StepEnd`](#)
events as well as handling [`CancelStepException`](#)s.
"""
function
runstep
(
stepfn
,
learner
,
phase
::
Phase
,
initialstate
=
(
;
)
)
state
=
PropDict
(
pairs
(
initialstate
)
)
handlefn
(
e
)
=
handle
(
learner
.
callbacks
.
runner
,
e
,
phase
,
learner
)
try
learner
.
step
=
state
handlefn
(
StepBegin
(
)
)
stepfn
(
handlefn
,
state
)
handlefn
(
StepEnd
(
)
)
return
state
catch
e
if
e
isa
CancelStepException
@
debug
"
Step skipped
"
error
=
e
else
rethrow
(
)
end
end
return
state
endUtilities
"""
fit!(learner, nepochs)
fit!(learner, nepochs, (trainiter, validiter))
Train `learner` for `nepochs` of training and validation each. Use data
iterators that are passed in. If none are given, use `learner.data.training`
and `learner.data.validation`.
## Examples
```julia
fit!(learner, 10)
fit!(learner, 10, (traindl, valdl))
```
"""
function
fit!
(
learner
,
nepochs
::
Int
,
(
trainiter
,
validiter
)
)
for
i
in
1
:
nepochs
epoch!
(
learner
,
TrainingPhase
(
)
,
trainiter
)
epoch!
(
learner
,
ValidationPhase
(
)
,
validiter
)
end
end
function
fit!
(
learner
,
nepochs
::
Int
)
fit!
(
learner
,
nepochs
,
(
learner
.
data
.
training
,
learner
.
data
.
validation
)
)
end