def main(step, args, model_state_dict, optimizer_state_dict):
#
# PART2
#
model = build_model(args).cuda()
one_ll = next(model.children()).weight
optimizer = FusedAdam(model.parameters())
ASP.init_model_for_pruning(model,
args.pattern,
verbosity=args.verbosity,
whitelist=args.whitelist,
allow_recompute_mask=args.allow_recompute_mask)
ASP.init_optimizer_for_pruning(optimizer)
torch.manual_seed(args.seed2)
model.load_state_dict(model_state_dict)
optimizer.load_state_dict(optimizer_state_dict)
print("Model sparsity is %s" %
("enabled" if ASP.sparsity_is_enabled() else "disabled"))
# train for a few steps with sparse weights
print("SPARSE :: ", one_ll)
step = train_loop(args, model, optimizer, step, args.num_sparse_steps_2)
def __init__(self, model_name="distilbert-base-uncased", num_labels=2, fused_opt=True, sparse=False):
super().__init__()
self.transformer = transformers.AutoModelForSequenceClassification.from_pretrained(model_name,
num_labels=num_labels,
output_attentions=False,
return_dict=True)
self.fused_opt = fused_opt
self.sparse = sparse
if self.fused_opt:
self.opt = apex.optimizers.FusedAdam(self.parameters(), lr=1e-5, adam_w_mode=True)
else:
self.opt = torch.optim.AdamW(self.parameters(), lr=1e-5)
if self.sparse:
print("Experimental: automatic sparsity enabled!")
print("For more info, see: https://github.com/NVIDIA/apex/tree/master/apex/contrib/sparsity")
self.transformer = self.transformer.cuda()
ASP.prune_trained_model(self.transformer, self.opt)
def main(args):
#
# PART1
#
torch.manual_seed(args.seed)
model = build_model(args).cuda()
one_ll = next(model.children()).weight
optimizer = FusedAdam(model.parameters())
ASP.init_model_for_pruning(model, args.pattern, whitelist=args.whitelist, allow_recompute_mask=args.allow_recompute_mask)
ASP.init_optimizer_for_pruning(optimizer)
step = 0
# train for a few steps with dense weights
print("DENSE :: ",one_ll)
step = train_loop(args, model, optimizer, step, args.num_dense_steps)
# simulate sparsity by inserting zeros into existing dense weights
ASP.compute_sparse_masks()
# train for a few steps with sparse weights
print("SPARSE :: ",one_ll)
step = train_loop(args, model, optimizer, step, args.num_sparse_steps)
#
# PART 2
#
torch.manual_seed(args.seed2)
# train for a few steps with sparse weights
print("SPARSE :: ",one_ll)
step = train_loop(args, model, optimizer, step, args.num_sparse_steps_2)
def main(args):
model = build_model(args).cuda()
one_ll = next(model.children()).weight
optimizer = FusedAdam(model.parameters())
# only prune linear layers, even though we also support conv1d, conv2d and conv3d
ASP.init_model_for_pruning(model,
"m4n2_1d",
whitelist=[torch.nn.Linear],
allow_recompute_mask=True)
ASP.init_optimizer_for_pruning(optimizer)
step = 0
# train for a few steps with dense weights
print("DENSE :: ", one_ll)
step = train_loop(args, model, optimizer, step, args.num_dense_steps)
# simulate sparsity by inserting zeros into existing dense weights
ASP.compute_sparse_masks()
# train for a few steps with sparse weights
print("SPARSE :: ", one_ll)
step = train_loop(args, model, optimizer, step, args.num_sparse_steps)
# recompute sparse masks
ASP.compute_sparse_masks()
# train for a few steps with sparse weights
print("SPARSE :: ", one_ll)
step = train_loop(args, model, optimizer, step, args.num_sparse_steps_2)
# turn off sparsity
print("SPARSE :: ", one_ll)
ASP.restore_pruned_weights()
# train for a few steps with dense weights
print("DENSE :: ", one_ll)
step = train_loop(args, model, optimizer, step, args.num_dense_steps_2)
def main(args):
#
# PART1
#
torch.manual_seed(args.seed)
model = build_model(args).cuda()
one_ll = next(model.children()).weight
optimizer = FusedAdam(model.parameters())
ASP.init_model_for_pruning(model,
args.pattern,
verbosity=args.verbosity,
whitelist=args.whitelist,
allow_recompute_mask=args.allow_recompute_mask)
ASP.init_optimizer_for_pruning(optimizer)
step = 0
# train for a few steps with dense weights
print("DENSE :: ", one_ll)
step = train_loop(args, model, optimizer, step, args.num_dense_steps)
# simulate sparsity by inserting zeros into existing dense weights
ASP.compute_sparse_masks()
# train for a few steps with sparse weights
print("SPARSE :: ", one_ll)
step = train_loop(args, model, optimizer, step, args.num_sparse_steps)
torch.save(
{
'step': step,
'verbosity': args.verbosity,
'seed2': args.seed2,
'pattern': args.pattern,
'whitelist': args.whitelist,
'allow_recompute_mask': args.allow_recompute_mask,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}, args.checkpoint_path)
def test_asp(self):
"""test Sparsity (ASP) and QAT toolkits together"""
try:
from apex.contrib.sparsity import ASP
except ImportError:
pytest.skip("ASP is not available.")
quant_modules.initialize()
model = LeNet()
quant_modules.deactivate()
optimizer = optim.SGD(model.parameters(), lr=0.01)
ASP.init_model_for_pruning(
model,
mask_calculator="m4n2_1d",
verbosity=2,
whitelist=[torch.nn.Linear, torch.nn.Conv2d, torch.nn.Conv3d, quant_nn.modules.quant_linear.QuantLinear],
allow_recompute_mask=False,
custom_layer_dict={
quant_nn.QuantConv1d: ['weight'],
quant_nn.QuantConv2d: ['weight'],
quant_nn.QuantConv3d: ['weight'],
quant_nn.QuantConvTranspose1d: ['weight'],
quant_nn.QuantConvTranspose2d: ['weight'],
quant_nn.QuantConvTranspose3d: ['weight'],
quant_nn.QuantLinear: ['weight']
})
ASP.init_optimizer_for_pruning(optimizer)
ASP.compute_sparse_masks()
model = model.to('cuda')
output = model(torch.empty(16, 1, 28, 28).to('cuda'))
optimizer.zero_grad()
loss = F.nll_loss(output, torch.randint(10, (16,), dtype=torch.int64))
loss.backward()
optimizer.step()