Skip to content

Fix an issue (#6) #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lyakaap wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Fix an issue (#6) #8

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4f0ed9f
Fix a issue (#6)
lyakaap Nov 20, 2018
3789520
Follow pycodestyle
lyakaap Nov 20, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 29 additions & 25 deletions torchstat/compute_memory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ def compute_memory(module, inp, out):
return compute_Pool2d_memory(module, inp, out)
else:
print(f"[Memory]: {type(module).__name__} is not supported!")
return (0, 0)
return 0, 0
pass


Expand All @@ -28,20 +28,21 @@ def num_params(module):

def compute_ReLU_memory(module, inp, out):
assert isinstance(module, (nn.ReLU, nn.ReLU6, nn.ELU, nn.LeakyReLU))
batch_size = inp.size()[0]
mread = batch_size * inp.size()[1:].numel()
mwrite = batch_size * inp.size()[1:].numel()

return (mread, mwrite)
mread = inp.numel()
mwrite = out.numel()

return mread, mwrite


def compute_PReLU_memory(module, inp, out):
assert isinstance(module, (nn.PReLU))
assert isinstance(module, nn.PReLU)

batch_size = inp.size()[0]
mread = batch_size * (inp.size()[1:].numel() + num_params(module))
mwrite = batch_size * inp.size()[1:].numel()
mread = batch_size * (inp[0].numel() + num_params(module))
mwrite = out.numel()

return (mread, mwrite)
return mread, mwrite


def compute_Conv2d_memory(module, inp, out):
Expand All @@ -50,39 +51,42 @@ def compute_Conv2d_memory(module, inp, out):
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

batch_size = inp.size()[0]
in_c = inp.size()[1]
out_c, out_h, out_w = out.size()[1:]

# This includes weighs with bias if the module contains it.
mread = batch_size * (inp.size()[1:].numel() + num_params(module))
mwrite = batch_size * out_c * out_h * out_w
return (mread, mwrite)
# This includes weights with bias if the module contains it.
mread = batch_size * (inp[0].numel() + num_params(module))
mwrite = out.numel()
return mread, mwrite


def compute_BatchNorm2d_memory(module, inp, out):
assert isinstance(module, nn.BatchNorm2d)
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())

batch_size, in_c, in_h, in_w = inp.size()
mread = batch_size * (inp[0].numel() + 2 * in_c)
mwrite = out.numel()

mread = batch_size * (inp.size()[1:].numel() + 2 * in_c)
mwrite = inp.size().numel()
return (mread, mwrite)
return mread, mwrite


def compute_Linear_memory(module, inp, out):
assert isinstance(module, nn.Linear)
assert len(inp.size()) == 2 and len(out.size()) == 2

batch_size = inp.size()[0]
mread = batch_size * (inp.size()[1:].numel() + num_params(module))
mwrite = out.size().numel()

return (mread, mwrite)
# This includes weights with bias if the module contains it.
mread = batch_size * (inp[0].numel() + num_params(module))
mwrite = out.numel()

return mread, mwrite


def compute_Pool2d_memory(module, inp, out):
assert isinstance(module, (nn.MaxPool2d, nn.AvgPool2d))
assert len(inp.size()) == 4 and len(inp.size()) == len(out.size())
batch_size = inp.size()[0]
mread = batch_size * inp.size()[1:].numel()
mwrite = batch_size * out.size()[1:].numel()
return (mread, mwrite)

mread = inp.numel()
mwrite = out.numel()

return mread, mwrite