save_group_sharded_model¶
使用save_group_sharded_model可以对group_sharded_parallel配置后的模型和优化器状态进行保存。
注解
此处需要注意，使用save_group_sharded_model保存模型，再次load时需要在调用group_sharded_parallel前对model和optimizer进行set_state_dict。
参数¶
model (Layer) - 使用group_sharded_parallel配置后的模型。
output (str) - 输出保存模型和优化器的文件夹路径。
optimizer (Optimizer，可选) - 使用group_sharded_parallel配置后的优化器，默认为None，表示不对优化器状态进行保存。
返回¶
无
代码示例¶
# required: distributed
import paddle
from paddle.fluid.dygraph.nn import Linear
from paddle.distributed import fleet
from paddle.distributed.sharding import group_sharded_parallel, save_group_sharded_model
fleet.init(is_collective=True)
group = paddle.distributed.new_group([0, 1])
model = Linear(1000, 1000)
clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=1.0)
optimizer = paddle.optimizer.AdamW(learning_rate=0.001, parameters=model.parameters(), weight_decay=0.00001, grad_clip=clip)
# wrap sharding model, optimizer and scaler
model, optimizer, scaler = group_sharded_parallel(model, optimizer, "p_g", scaler=scaler)
img, label = data
label.stop_gradient = True
img.stop_gradient = True
out = model(img)
loss = paddle.nn.functional.cross_entropy(input=out, label=label)
loss.backward()
optimizer.step()
optimizer.clear_grad()
# save model and optimizer state_dict
save_group_sharded_model(model, optimizer, output=output_dir)