Moved file to PhD project
This commit is contained in:
Eduardo Cueto Mendoza
parent
973b17a4e6
commit
8dab247c53
174
lenet.jl
174
lenet.jl
|
|
@ -1,174 +0,0 @@
|
|||
## Classification of MNIST dataset
|
||||
## with the convolutional neural network know as LeNet5.
|
||||
## This script also combines various
|
||||
## packages from the Julia ecosystem with Flux.
|
||||
using Flux
|
||||
using Flux.Data: DataLoader
|
||||
using Flux.Optimise: Optimiser, WeightDecay
|
||||
using Flux: onehotbatch, onecold, logitcrossentropy
|
||||
using Statistics, Random
|
||||
using Parameters: @with_kw
|
||||
using Logging: with_logger, global_logger
|
||||
using TensorBoardLogger: TBLogger, tb_overwrite, set_step!, set_step_increment!
|
||||
import ProgressMeter
|
||||
import MLDatasets
|
||||
import DrWatson: savename, struct2dict
|
||||
import BSON
|
||||
using CUDAapi
|
||||
|
||||
# LeNet5 "constructor".
|
||||
# The model can be adapted to any image size
|
||||
# and number of output classes.
|
||||
function LeNet5(; imgsize=(28,28,1), nclasses=10)
|
||||
out_conv_size = (imgsize[1]÷4 - 3, imgsize[2]÷4 - 3, 16)
|
||||
|
||||
return Chain(
|
||||
x -> reshape(x, imgsize..., :),
|
||||
Conv((5, 5), imgsize[end]=>6, relu),
|
||||
MaxPool((2, 2)),
|
||||
Conv((5, 5), 6=>16, relu),
|
||||
MaxPool((2, 2)),
|
||||
x -> reshape(x, :, size(x, 4)),
|
||||
Dense(prod(out_conv_size), 120, relu),
|
||||
Dense(120, 84, relu),
|
||||
Dense(84, nclasses)
|
||||
)
|
||||
end
|
||||
|
||||
function get_data(args)
|
||||
xtrain, ytrain = MLDatasets.MNIST.traindata(Float32, dir=args.datapath)
|
||||
xtest, ytest = MLDatasets.MNIST.testdata(Float32, dir=args.datapath)
|
||||
|
||||
xtrain = reshape(xtrain, 28, 28, 1, :)
|
||||
xtest = reshape(xtest, 28, 28, 1, :)
|
||||
|
||||
ytrain, ytest = onehotbatch(ytrain, 0:9), onehotbatch(ytest, 0:9)
|
||||
|
||||
train_loader = DataLoader(xtrain, ytrain, batchsize=args.batchsize, shuffle=true)
|
||||
test_loader = DataLoader(xtest, ytest, batchsize=args.batchsize)
|
||||
|
||||
return train_loader, test_loader
|
||||
end
|
||||
|
||||
loss(ŷ, y) = logitcrossentropy(ŷ, y)
|
||||
|
||||
function eval_loss_accuracy(loader, model, device)
|
||||
l = 0f0
|
||||
acc = 0
|
||||
ntot = 0
|
||||
for (x, y) in loader
|
||||
x, y = x |> device, y |> device
|
||||
ŷ = model(x)
|
||||
l += loss(ŷ, y) * size(x)[end]
|
||||
acc += sum(onecold(ŷ |> cpu) .== onecold(y |> cpu))
|
||||
ntot += size(x)[end]
|
||||
end
|
||||
return (loss = l/ntot |> round4, acc = acc/ntot*100 |> round4)
|
||||
end
|
||||
|
||||
## utility functions
|
||||
|
||||
num_params(model) = sum(length, Flux.params(model))
|
||||
|
||||
round4(x) = round(x, digits=4)
|
||||
|
||||
|
||||
# arguments for the `train` function
|
||||
@with_kw mutable struct Args
|
||||
η = 3e-4 # learning rate
|
||||
λ = 0 # L2 regularizer param, implemented as weight decay
|
||||
batchsize = 128 # batch size
|
||||
epochs = 20 # number of epochs
|
||||
seed = 0 # set seed > 0 for reproducibility
|
||||
cuda = true # if true use cuda (if available)
|
||||
infotime = 1 # report every `infotime` epochs
|
||||
checktime = 5 # Save the model every `checktime` epochs. Set to 0 for no checkpoints.
|
||||
tblogger = false # log training with tensorboard
|
||||
savepath = nothing # results path. If nothing, construct a default path from Args. If existing, may overwrite
|
||||
datapath = joinpath(homedir(), "Datasets", "MNIST") # data path: change to your data directory
|
||||
end
|
||||
|
||||
function train(; kws...)
|
||||
args = Args(; kws...)
|
||||
args.seed > 0 && Random.seed!(args.seed)
|
||||
use_cuda = args.cuda && CUDAapi.has_cuda_gpu()
|
||||
if use_cuda
|
||||
device = gpu
|
||||
@info "Training on GPU"
|
||||
else
|
||||
device = cpu
|
||||
@info "Training on CPU"
|
||||
end
|
||||
|
||||
## DATA
|
||||
train_loader, test_loader = get_data(args)
|
||||
@info "Dataset MNIST: $(train_loader.nobs) train and $(test_loader.nobs) test examples"
|
||||
|
||||
## MODEL AND OPTIMIZER
|
||||
model = LeNet5() |> device
|
||||
@info "LeNet5 model: $(num_params(model)) trainable params"
|
||||
|
||||
ps = Flux.params(model)
|
||||
|
||||
opt = ADAM(args.η)
|
||||
if args.λ > 0
|
||||
opt = Optimiser(opt, WeightDecay(args.λ))
|
||||
end
|
||||
|
||||
## LOGGING UTILITIES
|
||||
if args.savepath == nothing
|
||||
experiment_folder = savename("lenet", args, scientific=4,
|
||||
accesses=[:batchsize, :η, :seed, :λ]) # construct path from these fields
|
||||
args.savepath = joinpath("runs", experiment_folder)
|
||||
end
|
||||
if args.tblogger
|
||||
tblogger = TBLogger(args.savepath, tb_overwrite)
|
||||
set_step_increment!(tblogger, 0) # 0 auto increment since we manually set_step!
|
||||
@info "TensorBoard logging at \"$(args.savepath)\""
|
||||
end
|
||||
|
||||
function report(epoch)
|
||||
train = eval_loss_accuracy(train_loader, model, device)
|
||||
test = eval_loss_accuracy(test_loader, model, device)
|
||||
println("Epoch: $epoch Train: $(train) Test: $(test)")
|
||||
if args.tblogger
|
||||
set_step!(tblogger, epoch)
|
||||
with_logger(tblogger) do
|
||||
@info "train" loss=train.loss acc=train.acc
|
||||
@info "test" loss=test.loss acc=test.acc
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
## TRAINING
|
||||
@info "Start Training"
|
||||
report(0)
|
||||
for epoch in 1:args.epochs
|
||||
p = ProgressMeter.Progress(length(train_loader))
|
||||
|
||||
for (x, y) in train_loader
|
||||
x, y = x |> device, y |> device
|
||||
gs = Flux.gradient(ps) do
|
||||
ŷ = model(x)
|
||||
loss(ŷ, y)
|
||||
end
|
||||
Flux.Optimise.update!(opt, ps, gs)
|
||||
ProgressMeter.next!(p) # comment out for no progress bar
|
||||
end
|
||||
|
||||
epoch % args.infotime == 0 && report(epoch)
|
||||
if args.checktime > 0 && epoch % args.checktime == 0
|
||||
!ispath(args.savepath) && mkpath(args.savepath)
|
||||
modelpath = joinpath(args.savepath, "model.bson")
|
||||
let model=cpu(model), args=struct2dict(args)
|
||||
BSON.@save modelpath model epoch args
|
||||
end
|
||||
@info "Model saved in \"$(modelpath)\""
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
## Execution as a script
|
||||
if abspath(PROGRAM_FILE) == @__FILE__
|
||||
train()
|
||||
end
|
||||
Loading…
Reference in New Issue