在PyTorch权重矩阵上实施组套索

Question

我正在尝试在PyTorch中的神经网络的权重矩阵上实现组套索。

我已经编写了实现Group Lasso的代码，但是不确定这是否正确，确认或更正我的代码将非常有帮助。

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)

我希望torch.chunk函数将权重矩阵分解为小块，然后通过该块的L2规范和所有块之间的L1规范。