I am trying to implement Group Lasso on weight matrices of a neural network in PyTorch.
I have written the code to implement Group Lasso but am unsure if this is correct, confirmation or correction of my code will be very helpful.
def gl_norm(model, gl_lambda, num_blk):
gl_reg = torch.tensor(0., dtype=torch.float32).cuda()
for key in model:
for param in model[key].parameters():
dim = param.size()
if dim.__len__() > 1 and not model[key].skip_regularization:
div1 = list(torch.chunk(param,int(num_blk),1))
all_blks = []
for div2 in div1:
temp = list(torch.chunk(div2,int(num_blk),0))
for blk in temp:
all_blks.append(blk)
for l2_param in all_blks:
gl_reg += torch.norm(l2_param, 2)
return gl_reg * float(gl_lambda)
I expect the torch.chunk function to break up the weight matrix into small blocks which then go through L2 norm for the block and L1 norm between all the blocks.