在Flux.jl中保存训练过的模型的最佳方法？

Question

我有一个我训练过的模型，我想保存它以供将来使用并分发给其他人。使用Flux.jl保存经过训练的模型的最佳方法是什么

Answer 1

根据Flux.jl文档(https://fluxml.ai/Flux.jl/stable/saving/)，保存经过训练的模型的最佳方法是使用BSON.jl，方法是执行以下操作：

julia> using Flux

julia> model = Chain(Dense(10,5,relu),Dense(5,2),softmax)
Chain(Dense(10, 5, NNlib.relu), Dense(5, 2), NNlib.softmax)

julia> using BSON: @save

julia> @save "mymodel.bson" model

然后，您可以通过执行以下操作加载保存的模型：

julia> using Flux

julia> using BSON: @load

julia> @load "mymodel.bson" model

julia> model
Chain(Dense(10, 5, NNlib.relu), Dense(5, 2), NNlib.softmax)

Answer 2

如果你的模型没有像动态创建/调整大小的层这样的东西，你应该能够只保存权重而不是序列化整个模型。这可能比使用BSON.jl或Serialization stdlib来序列化整个模型(这两种方法都非常脆弱)要健壮得多。

权重可以通过weights=collect(params(cpu(model)))从模型中获得，并通过Flux.loadparams!(model, weights)加载回模型中。因此，人们只需要将数字数组的Vector保存到磁盘，而不是模型中更复杂的Julia端对象。因此，我建议使用如下模式：

function make_model(config)
   ...define layers, put them in a chain, etc...
    return model
end

# train model
...

# collect weights
weights=collect(params(cpu(model)))

# save them to disk somehow...

然后，当需要重新加载模型时，

weights = # load them from disk
fresh_model = make_model(config)
Flux.loadparams!(model, weights)

请注意，这种方法意味着您不能向make_model添加层并重新加载旧的权重；它们将不再是正确的大小。因此，您需要对代码和权重进行版本化，并确保它们匹配。

上周，我帮助制作了一个新的包LegolasFlux.jl来简化此模式(特别是，提供了一种使用箭头将权重以及您想要保存的任何其他配置参数、损失等保存到磁盘的方法)。它应该在两天内注册。