import numpy as np
import torch
from time import time

def dprint(s, DEBUG=True):
    if DEBUG:
        print(s)

def get_loss_and_acc(model, data_loader, criterion_sum):
    model.eval()

    device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
    loss = 0
    correct = 0
    
    with torch.no_grad():
        for data, target in data_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss += criterion_sum(output, target).item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    loss /= len(data_loader.dataset)
    acc = 100. * correct / len(data_loader.dataset)
    return loss, acc

def train(model, mask, train_loader, optimizer, criterion):
    ZERO_VAL = 1e-5
    t_start = time()
    
    model.train()
    device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
    loss = None

    for batch_idx, (data, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        data, labels = data.to(device), labels.to(device)
        output = model(data)
        loss = criterion(output, labels)
        loss.backward()

        for name, p in model.named_parameters():
            if 'weight' in name:
                grad = p.grad.data.cpu().numpy()
                p.grad.data = torch.from_numpy(grad * mask[name]).to(p.device)
        optimizer.step()
        
    dprint("\t\tEpoch time: {}".format(time()-t_start))

def make_mask(model):
    mask = {}
    for name, param in model.named_parameters():
        if 'weight' in name:
            tensor = param.data.cpu().numpy()
            mask[name] = np.ones_like(tensor)

    return mask

def permute_mask(mask):
    new_mask = {}
    for name in mask:
        new_mask[name] = np.random.choice(mask[name].flatten(), size=mask[name].shape, replace=False)
    return new_mask

def prune_by_percentile(model, mask, percent):
    for name, param in model.named_parameters():
        if 'weight' in name:
            weights = param.data.cpu().numpy()
            percentile_value = np.percentile(abs(weights[mask[name]!=0]), percent)

            new_mask = np.where(abs(weights) < percentile_value, 0, mask[name])
            param.data = torch.from_numpy(weights * new_mask).to(param.device)
            mask[name] = new_mask

def current_pruned_percent(model):
    total_params = 0
    total_active = 0

    for name, p in model.named_parameters():
        if 'weight' in name:
            weights = p.data.cpu().numpy()
            total_active += np.count_nonzero(weights)
            total_params += weights.size
    return total_active/total_params

def copy_params(model):
    params = {}
    for name, p in model.named_parameters():
        if 'weight' in name:
            params[name] = np.copy(p.data.cpu().numpy())
    return params

def extract_non_zero_params(weights):
    all_params = np.array([])
    for key in weights:
        layer = weights[key]
        all_params = np.concatenate((all_params, layer[np.nonzero(layer)]))
    return all_params

def reinitialize_model(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            param.data = torch.from_numpy(weights[name] * mask[name]).to(param.device)

def reinitialize_model_sample(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            new_weights = np.random.choice(weights, size=param.data.cpu().numpy().shape)
            param.data = torch.from_numpy(new_weights * mask[name]).to(param.device).float()

def copy_mask(mask):
    new_mask = {}
    for k in mask:
        new_mask[k] = mask[k].copy()
    return new_mask

def create_pruned_models(model, model_trainer, pruning_s, pruning_j, pruning_n):
    mask = make_mask(model)

    initial_params = copy_params(model)
    dprint("\tCreated initial model and mask")
    ret = {}

    ret[0] = (copy_mask(mask), initial_params)
    
    max_pruning = max(pruning_n)
    for n in range(1, max_pruning+1):
        dprint("\tStarting pruning n={}".format(n))
        for j in range(1,pruning_j+1):
            dprint("\t\tTraining epoch: {}".format(j))
            model_trainer(model, mask)
        dprint("\tPruning the model")
        prune_by_percentile(model, mask, pruning_s)
        
        if n in pruning_n:
            print("\tSaving pruned model")
            ret[n] = (copy(mask), copy_params(model))

        dprint("\tReinitializing masked model")
        reinitialize_model(model, initial_params, mask)

    return ret



7268021175cd9e028bef4d5d9b130dfb63a4d291
<mask token>


def train(model, mask, train_loader, optimizer, criterion):
    ZERO_VAL = 1e-05
    t_start = time()
    model.train()
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    loss = None
    for batch_idx, (data, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        data, labels = data.to(device), labels.to(device)
        output = model(data)
        loss = criterion(output, labels)
        loss.backward()
        for name, p in model.named_parameters():
            if 'weight' in name:
                grad = p.grad.data.cpu().numpy()
                p.grad.data = torch.from_numpy(grad * mask[name]).to(p.device)
        optimizer.step()
    dprint('\t\tEpoch time: {}'.format(time() - t_start))


<mask token>


def permute_mask(mask):
    new_mask = {}
    for name in mask:
        new_mask[name] = np.random.choice(mask[name].flatten(), size=mask[
            name].shape, replace=False)
    return new_mask


def prune_by_percentile(model, mask, percent):
    for name, param in model.named_parameters():
        if 'weight' in name:
            weights = param.data.cpu().numpy()
            percentile_value = np.percentile(abs(weights[mask[name] != 0]),
                percent)
            new_mask = np.where(abs(weights) < percentile_value, 0, mask[name])
            param.data = torch.from_numpy(weights * new_mask).to(param.device)
            mask[name] = new_mask


def current_pruned_percent(model):
    total_params = 0
    total_active = 0
    for name, p in model.named_parameters():
        if 'weight' in name:
            weights = p.data.cpu().numpy()
            total_active += np.count_nonzero(weights)
            total_params += weights.size
    return total_active / total_params


def copy_params(model):
    params = {}
    for name, p in model.named_parameters():
        if 'weight' in name:
            params[name] = np.copy(p.data.cpu().numpy())
    return params


<mask token>


def reinitialize_model(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            param.data = torch.from_numpy(weights[name] * mask[name]).to(param
                .device)


<mask token>

<mask token>


def dprint(s, DEBUG=True):
    if DEBUG:
        print(s)


<mask token>


def train(model, mask, train_loader, optimizer, criterion):
    ZERO_VAL = 1e-05
    t_start = time()
    model.train()
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    loss = None
    for batch_idx, (data, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        data, labels = data.to(device), labels.to(device)
        output = model(data)
        loss = criterion(output, labels)
        loss.backward()
        for name, p in model.named_parameters():
            if 'weight' in name:
                grad = p.grad.data.cpu().numpy()
                p.grad.data = torch.from_numpy(grad * mask[name]).to(p.device)
        optimizer.step()
    dprint('\t\tEpoch time: {}'.format(time() - t_start))


<mask token>


def permute_mask(mask):
    new_mask = {}
    for name in mask:
        new_mask[name] = np.random.choice(mask[name].flatten(), size=mask[
            name].shape, replace=False)
    return new_mask


def prune_by_percentile(model, mask, percent):
    for name, param in model.named_parameters():
        if 'weight' in name:
            weights = param.data.cpu().numpy()
            percentile_value = np.percentile(abs(weights[mask[name] != 0]),
                percent)
            new_mask = np.where(abs(weights) < percentile_value, 0, mask[name])
            param.data = torch.from_numpy(weights * new_mask).to(param.device)
            mask[name] = new_mask


def current_pruned_percent(model):
    total_params = 0
    total_active = 0
    for name, p in model.named_parameters():
        if 'weight' in name:
            weights = p.data.cpu().numpy()
            total_active += np.count_nonzero(weights)
            total_params += weights.size
    return total_active / total_params


def copy_params(model):
    params = {}
    for name, p in model.named_parameters():
        if 'weight' in name:
            params[name] = np.copy(p.data.cpu().numpy())
    return params


<mask token>


def reinitialize_model(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            param.data = torch.from_numpy(weights[name] * mask[name]).to(param
                .device)


def reinitialize_model_sample(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            new_weights = np.random.choice(weights, size=param.data.cpu().
                numpy().shape)
            param.data = torch.from_numpy(new_weights * mask[name]).to(param
                .device).float()


def copy_mask(mask):
    new_mask = {}
    for k in mask:
        new_mask[k] = mask[k].copy()
    return new_mask


<mask token>

<mask token>


def dprint(s, DEBUG=True):
    if DEBUG:
        print(s)


def get_loss_and_acc(model, data_loader, criterion_sum):
    model.eval()
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in data_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss += criterion_sum(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    loss /= len(data_loader.dataset)
    acc = 100.0 * correct / len(data_loader.dataset)
    return loss, acc


def train(model, mask, train_loader, optimizer, criterion):
    ZERO_VAL = 1e-05
    t_start = time()
    model.train()
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    loss = None
    for batch_idx, (data, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        data, labels = data.to(device), labels.to(device)
        output = model(data)
        loss = criterion(output, labels)
        loss.backward()
        for name, p in model.named_parameters():
            if 'weight' in name:
                grad = p.grad.data.cpu().numpy()
                p.grad.data = torch.from_numpy(grad * mask[name]).to(p.device)
        optimizer.step()
    dprint('\t\tEpoch time: {}'.format(time() - t_start))


<mask token>


def permute_mask(mask):
    new_mask = {}
    for name in mask:
        new_mask[name] = np.random.choice(mask[name].flatten(), size=mask[
            name].shape, replace=False)
    return new_mask


def prune_by_percentile(model, mask, percent):
    for name, param in model.named_parameters():
        if 'weight' in name:
            weights = param.data.cpu().numpy()
            percentile_value = np.percentile(abs(weights[mask[name] != 0]),
                percent)
            new_mask = np.where(abs(weights) < percentile_value, 0, mask[name])
            param.data = torch.from_numpy(weights * new_mask).to(param.device)
            mask[name] = new_mask


def current_pruned_percent(model):
    total_params = 0
    total_active = 0
    for name, p in model.named_parameters():
        if 'weight' in name:
            weights = p.data.cpu().numpy()
            total_active += np.count_nonzero(weights)
            total_params += weights.size
    return total_active / total_params


def copy_params(model):
    params = {}
    for name, p in model.named_parameters():
        if 'weight' in name:
            params[name] = np.copy(p.data.cpu().numpy())
    return params


<mask token>


def reinitialize_model(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            param.data = torch.from_numpy(weights[name] * mask[name]).to(param
                .device)


def reinitialize_model_sample(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            new_weights = np.random.choice(weights, size=param.data.cpu().
                numpy().shape)
            param.data = torch.from_numpy(new_weights * mask[name]).to(param
                .device).float()


def copy_mask(mask):
    new_mask = {}
    for k in mask:
        new_mask[k] = mask[k].copy()
    return new_mask


def create_pruned_models(model, model_trainer, pruning_s, pruning_j, pruning_n
    ):
    mask = make_mask(model)
    initial_params = copy_params(model)
    dprint('\tCreated initial model and mask')
    ret = {}
    ret[0] = copy_mask(mask), initial_params
    max_pruning = max(pruning_n)
    for n in range(1, max_pruning + 1):
        dprint('\tStarting pruning n={}'.format(n))
        for j in range(1, pruning_j + 1):
            dprint('\t\tTraining epoch: {}'.format(j))
            model_trainer(model, mask)
        dprint('\tPruning the model')
        prune_by_percentile(model, mask, pruning_s)
        if n in pruning_n:
            print('\tSaving pruned model')
            ret[n] = copy(mask), copy_params(model)
        dprint('\tReinitializing masked model')
        reinitialize_model(model, initial_params, mask)
    return ret

<mask token>


def dprint(s, DEBUG=True):
    if DEBUG:
        print(s)


def get_loss_and_acc(model, data_loader, criterion_sum):
    model.eval()
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in data_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss += criterion_sum(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    loss /= len(data_loader.dataset)
    acc = 100.0 * correct / len(data_loader.dataset)
    return loss, acc


def train(model, mask, train_loader, optimizer, criterion):
    ZERO_VAL = 1e-05
    t_start = time()
    model.train()
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    loss = None
    for batch_idx, (data, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        data, labels = data.to(device), labels.to(device)
        output = model(data)
        loss = criterion(output, labels)
        loss.backward()
        for name, p in model.named_parameters():
            if 'weight' in name:
                grad = p.grad.data.cpu().numpy()
                p.grad.data = torch.from_numpy(grad * mask[name]).to(p.device)
        optimizer.step()
    dprint('\t\tEpoch time: {}'.format(time() - t_start))


def make_mask(model):
    mask = {}
    for name, param in model.named_parameters():
        if 'weight' in name:
            tensor = param.data.cpu().numpy()
            mask[name] = np.ones_like(tensor)
    return mask


def permute_mask(mask):
    new_mask = {}
    for name in mask:
        new_mask[name] = np.random.choice(mask[name].flatten(), size=mask[
            name].shape, replace=False)
    return new_mask


def prune_by_percentile(model, mask, percent):
    for name, param in model.named_parameters():
        if 'weight' in name:
            weights = param.data.cpu().numpy()
            percentile_value = np.percentile(abs(weights[mask[name] != 0]),
                percent)
            new_mask = np.where(abs(weights) < percentile_value, 0, mask[name])
            param.data = torch.from_numpy(weights * new_mask).to(param.device)
            mask[name] = new_mask


def current_pruned_percent(model):
    total_params = 0
    total_active = 0
    for name, p in model.named_parameters():
        if 'weight' in name:
            weights = p.data.cpu().numpy()
            total_active += np.count_nonzero(weights)
            total_params += weights.size
    return total_active / total_params


def copy_params(model):
    params = {}
    for name, p in model.named_parameters():
        if 'weight' in name:
            params[name] = np.copy(p.data.cpu().numpy())
    return params


<mask token>


def reinitialize_model(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            param.data = torch.from_numpy(weights[name] * mask[name]).to(param
                .device)


def reinitialize_model_sample(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            new_weights = np.random.choice(weights, size=param.data.cpu().
                numpy().shape)
            param.data = torch.from_numpy(new_weights * mask[name]).to(param
                .device).float()


def copy_mask(mask):
    new_mask = {}
    for k in mask:
        new_mask[k] = mask[k].copy()
    return new_mask


def create_pruned_models(model, model_trainer, pruning_s, pruning_j, pruning_n
    ):
    mask = make_mask(model)
    initial_params = copy_params(model)
    dprint('\tCreated initial model and mask')
    ret = {}
    ret[0] = copy_mask(mask), initial_params
    max_pruning = max(pruning_n)
    for n in range(1, max_pruning + 1):
        dprint('\tStarting pruning n={}'.format(n))
        for j in range(1, pruning_j + 1):
            dprint('\t\tTraining epoch: {}'.format(j))
            model_trainer(model, mask)
        dprint('\tPruning the model')
        prune_by_percentile(model, mask, pruning_s)
        if n in pruning_n:
            print('\tSaving pruned model')
            ret[n] = copy(mask), copy_params(model)
        dprint('\tReinitializing masked model')
        reinitialize_model(model, initial_params, mask)
    return ret

import numpy as np
import torch
from time import time

def dprint(s, DEBUG=True):
    if DEBUG:
        print(s)

def get_loss_and_acc(model, data_loader, criterion_sum):
    model.eval()

    device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
    loss = 0
    correct = 0
    
    with torch.no_grad():
        for data, target in data_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss += criterion_sum(output, target).item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    loss /= len(data_loader.dataset)
    acc = 100. * correct / len(data_loader.dataset)
    return loss, acc

def train(model, mask, train_loader, optimizer, criterion):
    ZERO_VAL = 1e-5
    t_start = time()
    
    model.train()
    device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
    loss = None

    for batch_idx, (data, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        data, labels = data.to(device), labels.to(device)
        output = model(data)
        loss = criterion(output, labels)
        loss.backward()

        for name, p in model.named_parameters():
            if 'weight' in name:
                grad = p.grad.data.cpu().numpy()
                p.grad.data = torch.from_numpy(grad * mask[name]).to(p.device)
        optimizer.step()
        
    dprint("\t\tEpoch time: {}".format(time()-t_start))

def make_mask(model):
    mask = {}
    for name, param in model.named_parameters():
        if 'weight' in name:
            tensor = param.data.cpu().numpy()
            mask[name] = np.ones_like(tensor)

    return mask

def permute_mask(mask):
    new_mask = {}
    for name in mask:
        new_mask[name] = np.random.choice(mask[name].flatten(), size=mask[name].shape, replace=False)
    return new_mask

def prune_by_percentile(model, mask, percent):
    for name, param in model.named_parameters():
        if 'weight' in name:
            weights = param.data.cpu().numpy()
            percentile_value = np.percentile(abs(weights[mask[name]!=0]), percent)

            new_mask = np.where(abs(weights) < percentile_value, 0, mask[name])
            param.data = torch.from_numpy(weights * new_mask).to(param.device)
            mask[name] = new_mask

def current_pruned_percent(model):
    total_params = 0
    total_active = 0

    for name, p in model.named_parameters():
        if 'weight' in name:
            weights = p.data.cpu().numpy()
            total_active += np.count_nonzero(weights)
            total_params += weights.size
    return total_active/total_params

def copy_params(model):
    params = {}
    for name, p in model.named_parameters():
        if 'weight' in name:
            params[name] = np.copy(p.data.cpu().numpy())
    return params

def extract_non_zero_params(weights):
    all_params = np.array([])
    for key in weights:
        layer = weights[key]
        all_params = np.concatenate((all_params, layer[np.nonzero(layer)]))
    return all_params

def reinitialize_model(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            param.data = torch.from_numpy(weights[name] * mask[name]).to(param.device)

def reinitialize_model_sample(model, weights, mask):
    for name, param in model.named_parameters():
        if 'weight' in name:
            new_weights = np.random.choice(weights, size=param.data.cpu().numpy().shape)
            param.data = torch.from_numpy(new_weights * mask[name]).to(param.device).float()

def copy_mask(mask):
    new_mask = {}
    for k in mask:
        new_mask[k] = mask[k].copy()
    return new_mask

def create_pruned_models(model, model_trainer, pruning_s, pruning_j, pruning_n):
    mask = make_mask(model)

    initial_params = copy_params(model)
    dprint("\tCreated initial model and mask")
    ret = {}

    ret[0] = (copy_mask(mask), initial_params)
    
    max_pruning = max(pruning_n)
    for n in range(1, max_pruning+1):
        dprint("\tStarting pruning n={}".format(n))
        for j in range(1,pruning_j+1):
            dprint("\t\tTraining epoch: {}".format(j))
            model_trainer(model, mask)
        dprint("\tPruning the model")
        prune_by_percentile(model, mask, pruning_s)
        
        if n in pruning_n:
            print("\tSaving pruned model")
            ret[n] = (copy(mask), copy_params(model))

        dprint("\tReinitializing masked model")
        reinitialize_model(model, initial_params, mask)

    return ret


