MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / test_state_dict

Function test_state_dict

imperative/python/test/unit/module/test_module.py:378–398  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

376
377
378def test_state_dict():
379 data_shape = (2, 28)
380 data = tensor(np.random.random(data_shape))
381 mlp = MLP()
382 pred0 = mlp(data)
383
384 with BytesIO() as fout:
385 mge.save(mlp.state_dict(), fout)
386 fout.seek(0)
387 state_dict = mge.load(fout)
388 state_dict["extra"] = None
389 mlp1 = MLP()
390 mlp1.load_state_dict(state_dict, strict=False)
391 pred1 = mlp1(data)
392 np.testing.assert_allclose(pred0.numpy(), pred1.numpy(), atol=5e-6)
393 with pytest.raises(KeyError):
394 mlp1.load_state_dict(state_dict)
395 del state_dict["extra"]
396 del state_dict["dense0.bias"]
397 with pytest.raises(KeyError):
398 mlp1.load_state_dict(state_dict)
399
400
401def test_format_after_load_state_dict():

Callers

nothing calls this directly

Calls 7

MLPClass · 0.70
saveMethod · 0.45
state_dictMethod · 0.45
seekMethod · 0.45
loadMethod · 0.45
load_state_dictMethod · 0.45
numpyMethod · 0.45

Tested by

no test coverage detected