pytorch/caffe2/python/modeling/initializers_test.py
Simon Layton 2c3071fc4e Rework initializers to pass a class not object
Summary:
Changed tests
Moved to WeightInitializer, BiasInitializer keywords
Closes https://github.com/caffe2/caffe2/pull/682

Reviewed By: Yangqing

Differential Revision: D5138769

Pulled By: salexspb

fbshipit-source-id: 81d266100b2a95c64c0196c16670dfd34ea03e02
2017-05-30 09:06:56 -07:00

33 lines
1.1 KiB
Python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import unittest
from caffe2.python import brew, model_helper
from caffe2.python.modeling.initializers import Initializer
class InitializerTest(unittest.TestCase):
def test_fc_initializer(self):
model = model_helper.ModelHelper(name="test")
data = model.net.AddExternalInput("data")
fc1 = brew.fc(model, data, "fc1", dim_in=1, dim_out=1)
# no operator name set, will use default
fc2 = brew.fc(model, fc1, "fc2", dim_in=1, dim_out=1,
WeightInitializer=Initializer)
# no operator name set, will use custom
fc3 = brew.fc(model, fc2, "fc3", dim_in=1, dim_out=1,
WeightInitializer=Initializer,
weight_init=("ConstantFill", {}),
)
# operator name set, no initializer class set
fc4 = brew.fc(model, fc3, "fc4", dim_in=1, dim_out=1,
WeightInitializer=None,
weight_init=("ConstantFill", {})
)