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

Question

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

Answer 1

如果您的模型没有动态创建/调整大小的图层，您应该能够只保存权重而不是序列化整个模型。这比使用 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以使这种模式更容易（特别是提供了一种使用 Arrow 将权重以及您想要保存的任何其他配置参数、损失等保存到磁盘的方法）。两天后应该可以注册。

Answer 2

根据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)