torch-pruning

KO | EN

gitlite — search

#javascript #python #ai #typescript #go #rust #deep-learning #linux #macos #nodejs #vue #pytorch #openai #open-source

(★ 15)

Pruning methods for pytorch with an optimizer-like interface

torch-pruning Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# torch-pruning
Pruning methods for pytorch with an optimizer-like interface

## usage

```python
import torch
import pruning

net = # an arbitrary pytorch nn.Module instance
dataloader = # some pytorch dataloader instance

optimizer = torch.optim.SGD(net.parameters(), 0.01, weight_decay=1e-5)
# Init pruning method in the same way as optimizer
pruning = pruning.MagnitudePruning(net.parameters(), 0.1, local=True,
                                   exclude_biases=True)

# Save initial parameters for later
w_0 = pruning.clone_params()

def train(net, dataloader, n_epochs=1):
    # Some standard training loop ...
    for epoch in range(n_epochs):
        for x, y in dataloader:
            # Do actually set *pruned* weights to zero
            pruning.zero_params()
            y_hat = net(x)
            loss = criterion(y_hat, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

# Train epoch
train(net, dataloader, n_epochs=100)
# Do prune!
pruning.step()
# Rewind parameters to their values at init
pruning.rewind(w_0)
# Train the pruned model
train(net, dataloader, n_epochs=100)

# Do train and prune in a for loop, and then
# check if you have found a winning ticket
# ...

```

torch-pruning

torch-pruning Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

torch-pruning

torch-pruning Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README