MCPcopy Create free account
hub / github.com/Oneflow-Inc/oneflow / test_moduledict

Method test_moduledict

python/oneflow/test/modules/test_module.py:360–379  ·  view source on GitHub ↗
(test_case)

Source from the content-addressed store, hash-verified

358
359 @flow.unittest.skip_unless_1n1d()
360 def test_moduledict(test_case):
361 class ModuleDict(nn.Module):
362 def __init__(self):
363 super(ModuleDict, self).__init__()
364 self.choices = nn.ModuleDict(
365 {"conv": nn.Conv2d(10, 10, 3), "pool": nn.MaxPool2d(3)}
366 )
367 self.activations = nn.ModuleDict(
368 {"relu": nn.ReLU(), "prelu": nn.PReLU()}
369 )
370
371 def forward(self, x, choice, act):
372 x = self.choices[choice](x)
373 x = self.activations[act](x)
374 return x
375
376 model = ModuleDict()
377 input = flow.tensor(np.random.randn(4, 10, 32, 32), dtype=flow.float32)
378 output = model(input, "conv", "relu")
379 test_case.assertEqual(output.shape, flow.Size([4, 10, 30, 30]))
380
381 @flow.unittest.skip_unless_1n1d()
382 def test_module_submodule(test_case):

Callers

nothing calls this directly

Calls 2

ModuleDictClass · 0.70
SizeMethod · 0.45

Tested by

no test coverage detected