OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Fix BN for test phase

Browse Source
This commit is contained in:
Simon Layton 2016-09-12 15:28:01 -04:00
parent b861d9a264
commit 00c493864e
1 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
caffe2/python/cnn.py
View File

@ -470,9 +470,15 @@ class CNNModelHelper(object):
self.biases.append(bias)
blob_outs = [blob_out, blob_out + "_rm", blob_out + "_riv",
blob_out + "_sm", blob_out + "_siv"]
blob_outputs = self.net.SpatialBN(
[blob_in, scale, bias, blob_outs[1], blob_outs[2]], blob_outs,
order=self.order, **kwargs)
if kwargs['is_test']:
blob_outputs = self.net.SpatialBN(
[blob_in, scale, bias, blob_outs[1], blob_outs[2]], [blob_out],
order=self.order, **kwargs)
return blob_outputs
else:
blob_outputs = self.net.SpatialBN(
[blob_in, scale, bias, blob_outs[1], blob_outs[2]], blob_outs,
order=self.order, **kwargs)
# Return the output
return blob_outputs[0]