如何将 Tensor 保存下来并加载 ?

本文介绍了如何在PyTorch中使用torch.save和torch.load函数将Tensor保存为.pth文件并跨设备加载。示例展示了如何从CPU保存Tensor到磁盘,然后将其加载到GPU上,这对于在不同设备间转移模型或者处理大型Tensor非常实用。
摘要由CSDN通过智能技术生成

将 Tensor 保存为 pth 文件再加载即可。这种方法对于不同 device 的服务器之间的算子测试,以及保存 size 很大的 Tensor 非常方便。

a=torch.randn([3,4],device='cpu')
torch.save(a,"a.pth")    # 保存Tensor为pth文件
b=torch.load("a.pth",map_location="cuda:0")   # 指定加载的device
print(b)

输出:

tensor([[-0.2660,  0.5144,  0.3616,  1.5060],
        [-1.0543,  2.4541, -0.3537, -0.1083],
        [ 0.0604,  2.6085, -2.9730, -1.0179]], device='cuda:0')
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

ctrl A_ctrl C_ctrl V

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值