hpcaitech / TensorNVMe

A Python library transfers PyTorch tensors between CPU and NVMe
93 stars 18 forks source link

[feature] add async file writer #45

Closed ver217 closed 2 weeks ago

ver217 commented 2 weeks ago

Usage

from tensornvme.async_file_io import AsyncFileWriter
import torch.nn as nn
import torch
model = nn.Linear(10, 2)
with open("t.pkl", "wb") as f:
    f_writer.synchronize()
    torch.save(model.state_dict(), f_writer)
    f_writer.synchronize()
print(model.state_dict())
# output
# OrderedDict([('weight',
#               tensor([[-0.2586, -0.0774, -0.1293,  0.0518,  0.1004,  0.0637, -0.1227,  0.0488,
#                        -0.2836, -0.1508],
#                       [ 0.1144, -0.2542, -0.0783,  0.2590, -0.0972, -0.0133, -0.1492,  0.2579,
#                        -0.2547,  0.1043]])),
#              ('bias', tensor([0.0223, 0.1377]))])
print(torch.load("t.pkl"))
# output
# OrderedDict([('weight',
#              tensor([[-0.2586, -0.0774, -0.1293,  0.0518,  0.1004,  0.0637, -0.1227,  0.0488,
#                        -0.2836, -0.1508],
#                      [ 0.1144, -0.2542, -0.0783,  0.2590, -0.0972, -0.0133, -0.1492,  0.2579,
#                        -0.2547,  0.1043]])),
#              ('bias', tensor([0.0223, 0.1377]))])