def train_pp(model, loss_fn, optimizer, param, loader_train, loader_test, ratio_list, k=3, loader_val=None):
model.train()
ratio_count = 0
for epoch in range(param['num_epochs']):
print('Starting epoch %d / %d' % (epoch + 1, param['num_epochs']))
for t, (x, y) in enumerate(loader_train):
x_var, y_var = to_var(x), to_var(y.long())
scores = model(x_var)
loss = loss_fn(scores, y_var)
if (t + 1) % 100 == 0:
print('t = %d, loss = %.8f' % (t + 1, loss.data[0]))
optimizer.zero_grad()
loss.backward()
model.update_grad()
optimizer.step()
test(model, loader_test)
if epoch%k == 0 and epoch!=0 and ratio_count<len(ratio_list):
model.procedure_stb_pruning(ratio_list[ratio_count])
# model.procedure_weight_pruning(ratio_list[ratio_count])
ratio_count+=1
print(ratio_count, 'th','###########################pruning once######################################')
def train_rand(model, loss_fn, optimizer, param, loader_train, loader_test, ratio, loader_val=None):
model.train()
model.rand_mask(ratio)
for epoch in range(param['num_epochs']):
print('Starting epoch %d / %d' % (epoch + 1, param['num_epochs']))
for t, (x, y) in enumerate(loader_train):
x_var, y_var = to_var(x), to_var(y.long())
scores = model(x_var)
loss = loss_fn(scores, y_var)
if (t + 1) % 100 == 0:
print('t = %d, loss = %.8f' % (t + 1, loss.data[0]))
optimizer.zero_grad()
loss.backward()
optimizer.step()
test(model, loader_test)
def test(model, loader):
model.eval()
num_correct, num_samples = 0, len(loader.dataset)
for x, y in loader:
x_var = to_var(x, volatile=True)
scores = model(x_var)
_, preds = scores.data.cpu().max(1)
num_correct += (preds == y).sum()
acc = float(num_correct) / num_samples
print('Test accuracy: {:.2f}% ({}/{})'.format(
100. * acc,
num_correct,
num_samples,
))
return acc
def gen_mask(model,
loss_fn,
optimizer,
param,
loader_train,
loader_test,
ratio,
k=3,
loader_val=None):
test(model, loader_test)
model.train()
count = 0
ratio_ind = 0
for epoch in range(param['num_epochs']):
model.train()
print('Starting epoch %d / %d' % (epoch + 1, param['num_epochs']))
for t, (x, y) in enumerate(loader_train):
x_var, y_var = to_var(x), to_var(y.long())
scores = model(x_var)
loss = loss_fn(scores, y_var)
if (t + 1) % 100 == 0:
print('t = %d, loss = %.8f' % (t + 1, loss.item()))
optimizer.zero_grad()
loss.backward()
model.update_grad()
optimizer.step()
# print(epoch,t)
# test(model, loader_test)
if (epoch + 1) % k == 0 and ratio_ind < len(ratio):
print(
' pruning some filters which are in convolution layes , pruning ratio:%.3f'
% ratio[ratio_ind])
if ratio_ind == 0:
model.com_mask2(ratio[ratio_ind], 0)
else:
model.com_mask2(ratio[ratio_ind], ratio[ratio_ind - 1])
model.set_masks(model.mask)
model.zero_accmgrad()
ratio_ind += 1
else:
model.set_masks(model.mask)
prune_rate(model)
print('modify learning rate')
lr = param['learning_rate'] * (0.5**((epoch - k * len(ratio)) // 30))
# lr = param['learning_rate'] * (0.5 ** ((epoch - 1) // 30))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
print('epoch', epoch)
test(model, loader_test)
count += 1
def set_mask(self, mask):
self.mask = to_var(mask, requires_grad=False)
self.weight.data = self.weight.data.cuda() * self.mask.data
self.mask_flag = True
def gen_mask(model,
loss_fn,
optimizer,
param,
loader_train,
loader_test,
ratio,
k=3,
loader_val=None):
model.train()
count = 0
ratio_ind = 0
for epoch in range(param['num_epochs']):
model.train()
print('Starting epoch %d / %d' % (epoch + 1, param['num_epochs']))
for t, (x, y) in enumerate(loader_train):
# print(t)
x_var, y_var = to_var(x), to_var(y.long())
# x_var = x.cuda(); y_var = y.long().cuda()
scores = model(x_var)
loss = loss_fn(scores, y_var)
if (t + 1) % 100 == 0:
print('t = %d, loss = %.8f' %
(t + 1, loss.item())) #loss.data[0]))
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (epoch + 1) % k == 0 and ratio_ind < len(ratio):
print(
' pruning some filters which are in convolution layes , pruning ratio:%.3f'
% ratio[ratio_ind])
if ratio_ind == 0:
model.com_mask3(ratio[ratio_ind], 0)
model.set_masks(model.mask)
print(model.layer_list)
print(model.filters_array)
print(len(model.filters_array))
print('not prune')
# model = prune_vgg(model, model.layer_list, model.filters_array)
# init_model(model)
# optimizer.params = model.parameters()
else:
model.com_mask3(ratio[ratio_ind], ratio[ratio_ind - 1])
model.set_masks(model.mask)
print(model.layer_list)
print(model.filters_array)
print(len(model.filters_array))
print('not prune')
ratio_ind += 1
# 当所有的剪枝率满足之后，再训练k个epochs，然后再剪枝——因为比如剪到80的时候，会陷入10%的怪圈，见同where中的结果
if epoch > k * (len(ratio) + 1):
print(model.layer_list)
print(model.filters_array)
print(len(model.filters_array))
print('real pruning')
model = prune_vgg(model, model.layer_list, model.filters_array)
return model
print('modify learning rate')
lr = param['learning_rate'] * (0.5**((epoch - k * len(ratio)) // 100))
# lr = param['learning_rate'] * (0.5 ** ((epoch - 1) // 30))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
print('epoch', epoch)
test(model, loader_test)
count += 1
return model
def get_mask(self):
# print(self.mask_flag)
return to_var(self.mask, requires_grad=False)
def get_mask(self):
# if self.mask_flag == False :
# return None
# else :
return to_var(self.mask, requires_grad=False)
