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

Jump to bottom

Freeze pruned weights method not efficient #10

Open

guoyuntu opened this issue Mar 11, 2021 · 1 comment

guoyuntu commented Mar 11, 2021

In 'main.py' line 257 - 262, the author used the following codes to freeze the pruned weights:

for name, p in model.named_parameters():
        if 'weight' in name:
            tensor = p.data.cpu().numpy()
            grad_tensor = p.grad.data.cpu().numpy()
            grad_tensor = np.where(tensor < EPS, 0, grad_tensor)
            p.grad.data = torch.from_numpy(grad_tensor).to(device)

which causes a heavy burden for CPU2GPU I/O.
I will recommend conducting the freezing operation on GPU directly, the following codes helps:

    for name, p in model.named_parameters():
        if 'weight' in name:
            tensor = p.data
            grad_tensor = p.grad
            grad_tensor = torch.where(tensor.abs() < EPS, torch.zeros_like(grad_tensor), grad_tensor)
            p.grad.data = grad_tensor

The text was updated successfully, but these errors were encountered:

guoyuntu changed the title ~~Freeze pruned weights method not effective~~ Freeze pruned weights method not efficient

bainro commented Dec 24, 2021

A batch size of 200 on mnist + lenet5 went from 12 seconds per epoch to 6 with your changes. Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment