将 PyTorch 模型 state_dict 保存到 redis 缓存中

Question

我正在构建一个分布式参数/服务器类型架构，并希望通过 Azure 上的表解决方案来传达 model 更新。

我很难找到有关将 PyTorch 模型 state_dict 保存到 redis 缓存中的任何有用信息。 I've given up on Azure Cosmos tables because of the size limit (64kb) per entity and looked toward redis since model state_dict params/weights are much larger, even for a small model.

有人对我如何追求这个有任何建议吗？ 或者如果这甚至是可能的？

Answer 1

我的解决方案（在@GuyKorland 上面评论之后）是 RedisAI。

我为 model 数据实现了键值机制，并以这种方式在 VM 之间进行通信。

for name, param in model.named_parameters():
    redisai_client.tensorset(f'{name}',param.data.numpy().cpu().detach())