Hi,
I want to add group lasso regularization to cross entropy loss. The starter code can be seen here: https://github.com/pytorch/examples/blob/master/word_language_model/main.py#L158
Group lasso regularization can be viewed as a function of model.rnn.weight_ih. To add group lasso, I modify this part of code from
criterion = nn.CrossEntropyLoss()
loss = criterion(output.view(-1, ntokens), targets)
loss.backward()
to
def add_dimension_glasso(var, dim=0):
return var.pow(2).sum(dim=dim).add(1e-8).pow(1/2.).sum()
criterion = nn.CrossEntropyLoss()
loss = criterion(output.view(-1, ntokens), targets)
weight_l0 = torch.cat((model.rnn.weight_ih_l0, model.rnn.weight_hh_l0), 1) # shape (4*hidden_size x input_size)
weight_l1 = torch.cat((model.rnn.weight_ih_l1, model.rnn.weight_hh_l1), 1)
weight_l2 = model.decoder.weight # decoder layer weight of shape (out_features x in_features)
dimension_glasso_reg = 0.003 * add_dimension_glasso(weight_l0, 0) + \
0.003 * add_dimension_glasso(weight_l1, 0) + \
0.003 * add_dimension_glasso(weight_l0, 1) + \
0.003 * add_dimension_glasso(weight_l1, 1)
structure_glasso_reg = 0.00245 * add_structure_glasso(weight_l0, weight_l1, 2) + \
0.00245 * add_structure_glasso(weight_l1, weight_l2, 1)
loss += dimension_glasso_reg
loss += structure_glasso_reg
loss.backward()
Does it work? Thank you.