gryf/coach
1
0
Fork 0
mirror of https://github.com/gryf/coach.git synced 2026-01-06 13:54:21 +01:00
Code Issues Projects Releases Wiki Activity

Added Custom Initialisation for MXNet Heads (#86)

Browse Source 
* Added NormalizedRSSInitializer, using same method as TensorFlow backend, but changed name since ‘columns’ have different meaning in dense layer weight matrix in MXNet.

* Added unit test for NormalizedRSSInitializer.
This commit is contained in:
Thom Lane
2018-11-16 08:15:43 -08:00
committed by Scott Leishman
parent 101c55d37d
commit 81bac050d7
5 changed files with 59 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
rl_coach/architectures/mxnet_components/heads/v_head.py
View File

@@ -4,7 +4,8 @@ from types import ModuleType
import mxnet as mx
from mxnet.gluon.loss import Loss, HuberLoss, L2Loss
from mxnet.gluon import nn
from rl_coach.architectures.mxnet_components.heads.head import Head, HeadLoss, LossInputSchema
from rl_coach.architectures.mxnet_components.heads.head import Head, HeadLoss, LossInputSchema,\
NormalizedRSSInitializer
from rl_coach.architectures.mxnet_components.heads.head import LOSS_OUT_TYPE_LOSS
from rl_coach.base_parameters import AgentParameters
from rl_coach.core_types import VStateValue
@@ -79,7 +80,7 @@ class VHead(Head):
self.loss_type = loss_type
self.return_type = VStateValue
with self.name_scope():
self.dense = nn.Dense(units=1)
self.dense = nn.Dense(units=1, weight_initializer=NormalizedRSSInitializer(1.0))
def loss(self) -> Loss:
"""