gryf/coach
1
0
Fork 0
mirror of https://github.com/gryf/coach.git synced 2025-12-18 03:30:19 +01:00
Code Issues Projects Releases Wiki Activity

Channel order transpose, for image embedder. Updated unit test. (#87)

Browse Source
This commit is contained in:
Thom Lane
2018-11-19 05:39:03 -08:00
committed by Gal Novik
parent ff816b347d
commit 7ba1a4393f
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
rl_coach/architectures/mxnet_components/embedders/image_embedder.py
View File

@@ -70,7 +70,6 @@ class ImageEmbedder(InputEmbedder):
:param x: image representing environment state, of shape (batch_size, in_channels, height, width).
:return: embedding of environment state, of shape (batch_size, channels).
"""
if len(x.shape) != 4 and self.scheme != EmbedderScheme.Empty:
raise ValueError("Image embedders expect the input size to have 4 dimensions. The given size is: {}"
.format(x.shape))
# convert from NHWC to NCHW (default for MXNet Convolutions)
x = x.transpose((0,3,1,2))
return super(ImageEmbedder, self).hybrid_forward(F, x, *args, **kwargs)

3
rl_coach/tests/architectures/mxnet_components/embedders/test_image_embedder.py
View File

@@ -15,7 +15,8 @@ def test_image_embedder():
params = InputEmbedderParameters(scheme=EmbedderScheme.Medium)
emb = ImageEmbedder(params=params)
emb.initialize()
input_data = mx.nd.random.uniform(low=0, high=1, shape=(10, 3, 244, 244))
# input is NHWC, and not MXNet default NCHW
input_data = mx.nd.random.uniform(low=0, high=1, shape=(10, 244, 244, 3))
output = emb(input_data)
assert len(output.shape) == 2 # since last block was flatten
assert output.shape[0] == 10 # since batch_size is 10