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)

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

2 个答案: