def test():
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
correct += pred.eq(target.data.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')
return accuracy
if args.pretrained:
model.load_state_dict(torch.load('saves/elt_0.0_0.pth'))
accuracy = test()
# Initial training
print("--- Initial training ---")
train(args.epochs, decay=args.decay, threshold=0.0)
accuracy = test()
torch.save(model.state_dict(), 'saves/elt_'+str(args.decay)+'_'+str(args.reg_type)+'.pth')
util.log(args.log, f"initial_accuracy {accuracy}")
#util.print_nonzeros(model)
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
correct += pred.eq(target.data.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
#print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')
return accuracy
model.load_state_dict(torch.load(args.model+'.pth'))
# Initial training
print("--- Pruning ---")
for name, p in model.named_parameters():
if 'mask' in name:
continue
tensor = p.data.cpu().numpy()
threshold = args.sensitivity*np.std(tensor)
print(f'Pruning with threshold : {threshold} for layer {name}')
new_mask = np.where(abs(tensor) < threshold, 0, tensor)
p.data = torch.from_numpy(new_mask).to(device)
accuracy = test()
util.print_nonzeros(model)
print("--- Finetuning ---")