def test_release_blobs_when_used(self):
m = model_helper.ModelHelper()
fc1 = brew.fc(m, "data", "x", dim_in=2, dim_out=2)
fc2 = brew.fc(m, fc1, "y", dim_in=2, dim_out=2)
fc3 = brew.fc(m, fc1, "z", dim_in=2, dim_out=2)
fc4 = brew.fc(m, fc2, "u", dim_in=2, dim_out=2)
m.net.Alias(["u"], ["u_alias"])
brew.sum(m, [fc3, fc4], "out")
with_frees = memonger.release_blobs_when_used(m.net.Proto(), set("data"))
expect_frees = {"x", "y", "z"} # out is external output
# and u is aliased so cannot be freed
found_frees = set()
for op in with_frees.op:
if op.type == "Free":
self.assertFalse(op.input[0] in found_frees) # no double frees
found_frees.add(op.input[0])
else:
# Check a freed blob is not used anymore
for inp in op.input:
self.assertFalse(inp in found_frees)
for outp in op.output:
self.assertFalse(outp in found_frees)
self.assertEqual(expect_frees, found_frees)
def test_release_blobs_when_used(self):
m = model_helper.ModelHelper()
fc1 = brew.fc(m, "data", "x", dim_in=2, dim_out=2)
fc2 = brew.fc(m, fc1, "y", dim_in=2, dim_out=2)
fc3 = brew.fc(m, fc1, "z", dim_in=2, dim_out=2)
fc4 = brew.fc(m, fc2, "u", dim_in=2, dim_out=2)
m.net.Alias(["u"], ["u_alias"])
brew.sum(m, [fc3, fc4], "out")
with_frees = memonger.release_blobs_when_used(m.net.Proto(),
set("data"))
expect_frees = {"x", "y", "z"} # out is external output
# and u is aliased so cannot be freed
found_frees = set()
for op in with_frees.op:
if op.type == "Free":
self.assertFalse(op.input[0] in found_frees) # no double frees
found_frees.add(op.input[0])
else:
# Check a freed blob is not used anymore
for inp in op.input:
self.assertFalse(inp in found_frees)
for outp in op.output:
self.assertFalse(outp in found_frees)
self.assertEqual(expect_frees, found_frees)
def optimize_memory(model, gpu_id=0, blobs_to_keep=[]):
"""Save GPU memory through releasing blobs."""
blobs_to_keep_all_devices = set()
if blobs_to_keep is not None:
for blob_name in blobs_to_keep:
blobs_to_keep_all_devices.add("gpu_{}/{}".format(
gpu_id, blob_name))
model.net._net = memonger.release_blobs_when_used(
model.net.Proto(), blobs_to_keep_all_devices)