def test_av_load_multiple_batches_per_layer(self) -> None:
def save_load_and_assert_batch(layer_path, total_num_batches, batch,
n_batch_name):
# save n-th batch and verify the number of saved batches
AV.save(
tmpdir,
model_id,
DEFAULT_IDENTIFIER,
"layer1.0.conv1",
batch,
n_batch_name,
)
loaded_dataset = AV.load(tmpdir, model_id, DEFAULT_IDENTIFIER,
"layer1.0.conv1", n_batch_name)
assertTensorAlmostEqual(self, next(iter(loaded_dataset)), batch,
0.0)
loaded_dataset_for_layer = AV.load(tmpdir, model_id,
DEFAULT_IDENTIFIER,
"layer1.0.conv1")
self.assertEqual(
loaded_dataset_for_layer.__len__(),
total_num_batches,
)
with tempfile.TemporaryDirectory() as tmpdir:
b0 = torch.randn(64, 16)
b1 = torch.randn(64, 16)
b2 = torch.randn(64, 16)
model_id = "dummy"
model_path = AV._assemble_model_dir(tmpdir, model_id)
layer_path = AV._assemble_file_path(model_path, DEFAULT_IDENTIFIER,
"layer1.0.conv1")
# save first batch and verify the number of saved batches
save_load_and_assert_batch(layer_path, 1, b0, "0")
# save second batch and verify the number of saved batches
save_load_and_assert_batch(layer_path, 2, b1, "1")
# save third batch and verify the number of saved batches
save_load_and_assert_batch(layer_path, 3, b2, "2")
def test_av_save_multiple_batches_per_layer(self) -> None:
def save_and_assert_batch(layer_path, total_num_batches, batch,
n_batch_name):
# save n-th batch and verify the number of saved batches
AV.save(
tmpdir,
model_id,
DEFAULT_IDENTIFIER,
"layer1.0.conv1",
batch,
n_batch_name,
)
self.assertEqual(
len(glob.glob("/".join([layer_path, "*.pt"]))),
total_num_batches,
)
self.assertTrue(
AV.exists(tmpdir, model_id, DEFAULT_IDENTIFIER,
"layer1.0.conv1", n_batch_name))
with tempfile.TemporaryDirectory() as tmpdir:
b0 = torch.randn(64, 16)
b1 = torch.randn(64, 16)
b2 = torch.randn(64, 16)
model_id = "dummy"
model_path = AV._assemble_model_dir(tmpdir, model_id)
layer_path = AV._assemble_file_path(model_path, DEFAULT_IDENTIFIER,
"layer1.0.conv1")
# save first batch and verify the number of saved batches
save_and_assert_batch(layer_path, 1, b0, "0")
# save second batch and verify the number of saved batches
save_and_assert_batch(layer_path, 2, b1, "1")
# save third batch and verify the number of saved batches
save_and_assert_batch(layer_path, 3, b2, "2")
def test_av_save_multi_layer(self) -> None:
with tempfile.TemporaryDirectory() as tmpdir:
av_0 = torch.randn(64, 16)
av_1 = torch.randn(64, 16)
av_2 = torch.randn(64, 16)
model_path = AV._assemble_model_dir(tmpdir, "dummy")
# save first layer
AV.save(tmpdir, "dummy", DEFAULT_IDENTIFIER, "layer1.0.conv1",
av_0, "0")
self.assertEqual(len(glob.glob(model_path + "*")), 1)
# add two new layers at once
AV.save(
tmpdir,
"dummy",
DEFAULT_IDENTIFIER,
["layer1.0.conv2", "layer1.1.conv1"],
[av_1, av_2],
"0",
)
self.assertEqual(len(glob.glob(model_path + "/*/*/*")), 3)
# overwrite the first saved layer
AV.save(tmpdir, "dummy", DEFAULT_IDENTIFIER, "layer1.0.conv1",
av_0, "0")
self.assertEqual(len(glob.glob(model_path + "/*/*/*")), 3)
# save a new version of the first layer
idf1 = str(int(datetime.now().microsecond))
self.assertFalse(AV.exists(tmpdir, "dummy", idf1,
"layer1.0.conv1"))
AV.save(tmpdir, "dummy", idf1, "layer1.0.conv1", av_0, "0")
self.assertTrue(AV.exists(tmpdir, "dummy", idf1, "layer1.0.conv1"))
self.assertEqual(len(glob.glob(model_path + "/*/*/*")), 4)