Allow Dendritic MLP to have multiple output heads #527
Conversation
akashvelu
requested review from
subutai and
karangrewal
and removed request for
subutai
There was a problem hiding this comment.
I left some minor comments. Before I approve, can you explain the rationale behind having multiple output heads and how you plan to use them? Is it so that you can index the one corresponding to the certain task during training, so that all other heads are unmodified?
| dendrite_weight_sparsity=0.95, | ||
| weight_sparsity=0.95, weight_init="modified", dendrite_init="modified", | ||
| freeze_dendrites=False, output_nonlinearity=None, | ||
| dendritic_layer_class=AbsoluteMaxGatingDendriticLayer, |
There was a problem hiding this comment.
I think the indentation here is fine, esp since that's how this stack overlow article does it.
Also, I copied this style from the model utils file.
There was a problem hiding this comment.
I left some minor comments. Before I approve, can you explain the rationale behind having multiple output heads and how you plan to use them? Is it so that you can index the one corresponding to the certain task during training, so that all other heads are unmodified?
Multiple output heads are used by the policy network, which outputs the mean and diagonal log std, each of which are separate heads, of a gaussian.
| self._output_layers = nn.ModuleList() | ||
| for out_size in output_size: | ||
| output_layer = nn.Sequential() | ||
| output_linear = SparseWeights(module=nn.Linear(input_size, out_size), |
There was a problem hiding this comment.
shouldn't the size of the output layer be (hidden_sizes[-1], out_size)?
There was a problem hiding this comment.
input_size is updated in line 163 to be the last hidden size, so it should be the same.
| self._output_layer.add_module("output_linear", output_linear) | ||
| self._single_output_head = not isinstance(output_size, Iterable) | ||
| if self._single_output_head: | ||
| output_size = (output_size,) |
There was a problem hiding this comment.
Can you update the docstring regarding the format of output_size?
These changes modify the dendritic_mlp class to have multiple output heads which all use the same hidden layer output.