def test_tensor_serialization():
with TemporaryFile() as f:
data = np.random.randint(low=0, high=7, size=[233])
a = Tensor(data, device="cpu0", dtype=np.int32)
mge.save(a, f)
f.seek(0)
b = mge.load(f)
np.testing.assert_equal(a.numpy(), data)
assert b.device.logical_name == "cpu0:0"
assert b.dtype == np.int32
with TemporaryFile() as f:
a = Parameter(np.random.random(size=(233, 2)).astype(np.float32))
mge.save(a, f)
f.seek(0)
b = mge.load(f)
assert isinstance(b, Parameter)
np.testing.assert_equal(a.numpy(), b.numpy())
with TemporaryFile() as f:
a = Tensor(np.random.random(size=(2, 233)).astype(np.float32))
mge.save(a, f)
f.seek(0)
b = mge.load(f)
assert type(b) is Tensor
np.testing.assert_equal(a.numpy(), b.numpy())
with TemporaryFile() as f:
a = Tensor(np.random.random(size=(2, 233)).astype(np.float32))
mge.save(a, f)
f.seek(0)
b = mge.load(f, map_location="cpux")
assert type(b) is Tensor
assert "cpu" in str(b.device)
np.testing.assert_equal(a.numpy(), b.numpy())
with TemporaryFile() as f:
if mge.is_cuda_available():
device_org = mge.get_default_device()
mge.set_default_device("gpu0")
a = Tensor(np.random.random(size=(2, 233)).astype(np.float32))
mge.save(a, f)
f.seek(0)
mge.set_default_device("cpux")
b = mge.load(f, map_location={"gpu0": "cpu0"})
assert type(b) is Tensor
assert "cpu0" in str(b.device)
np.testing.assert_equal(a.numpy(), b.numpy())
mge.set_default_device(device_org)
with TemporaryFile() as f:
a = Tensor(0)
a.qparams.scale = Tensor(1.0)
mge.save(a, f)
f.seek(0)
b = mge.load(f)
assert isinstance(b.qparams.scale, Tensor)
np.testing.assert_equal(b.qparams.scale.numpy(), 1.0)
def test_tensor_serialization():
def tensor_eq(a, b):
assert a.dtype == b.dtype
assert a.device == b.device
np.testing.assert_equal(a.numpy(), b.numpy())
with TemporaryFile() as f:
data = np.random.randint(low=0, high=7, size=[233])
a = Tensor(data, device="xpux", dtype=np.int32)
pickle.dump(a, f)
f.seek(0)
b = pickle.load(f)
np.testing.assert_equal(a.numpy(), b.numpy())
with TemporaryFile() as f:
a = Parameter(np.random.random(size=(233, 2)).astype(np.float32))
pickle.dump(a, f)
f.seek(0)
b = pickle.load(f)
assert isinstance(b, Parameter)
np.testing.assert_equal(a.numpy(), b.numpy())
with TemporaryFile() as f:
a = Tensor(np.random.random(size=(2, 233)).astype(np.float32))
pickle.dump(a, f)
f.seek(0)
b = pickle.load(f)
assert type(b) is Tensor
np.testing.assert_equal(a.numpy(), b.numpy())
with TemporaryFile() as f:
a = Tensor(np.random.random(size=(2, 233)).astype(np.float32))
mge.save(a, f)
f.seek(0)
b = mge.load(f, map_location="cpux")
assert type(b) is Tensor
assert "cpu" in str(b.device)
np.testing.assert_equal(a.numpy(), b.numpy())
with TemporaryFile() as f:
if mge.is_cuda_available():
device_org = mge.get_default_device()
mge.set_default_device("gpu0")
a = Tensor(np.random.random(size=(2, 233)).astype(np.float32))
mge.save(a, f)
f.seek(0)
mge.set_default_device("cpux")
b = mge.load(f, map_location={"gpu0": "cpu0"})
assert type(b) is Tensor
assert "cpu0" in str(b.device)
np.testing.assert_equal(a.numpy(), b.numpy())
mge.set_default_device(device_org)
def test_save_load():
net = Simple()
optim = optimizer.SGD(net.parameters(), lr=1.0, momentum=0.9)
optim.clear_grad()
gm = ad.GradManager().attach(net.parameters())
data = tensor([2.34])
with gm:
loss = net(data)
gm.backward(loss)
optim.step()
model_name = "simple.pkl"
print("save to {}".format(model_name))
mge.save(
{
"name": "simple",
"state_dict": net.state_dict(),
"opt_state": optim.state_dict(),
},
model_name,
)
# Load param to cpu
checkpoint = mge.load(model_name, map_location="cpu0")
device_save = mge.get_default_device()
mge.set_default_device("cpu0")
net = Simple()
net.load_state_dict(checkpoint["state_dict"])
optim = optimizer.SGD(net.parameters(), lr=1.0, momentum=0.9)
optim.load_state_dict(checkpoint["opt_state"])
print("load done")
os.remove("simple.pkl")
with gm:
loss = net([1.23])
gm.backward(loss)
optim.step()
# Restore device
mge.set_default_device(device_save)