def _get_model_and_mutators():
base_model = Net()
script_module = torch.jit.script(base_model)
base_model_ir = convert_to_graph(script_module, base_model)
base_model_ir.evaluator = DebugEvaluator()
mutators = process_inline_mutation(base_model_ir)
return base_model_ir, mutators
def test_input_choice(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 3, kernel_size=1)
self.conv2 = nn.Conv2d(3, 5, kernel_size=1)
self.input = nn.InputChoice(2)
def forward(self, x):
x1 = self.conv1(x)
x2 = self.conv2(x)
return self.input([x1, x2])
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutator = mutators[0].bind_sampler(EnumerateSampler())
model1 = mutator.apply(model)
model2 = mutator.apply(model)
self.assertEqual(
self._get_converted_pytorch_model(model1)(torch.randn(1, 3, 3,
3)).size(),
torch.Size([1, 3, 3, 3]))
self.assertEqual(
self._get_converted_pytorch_model(model2)(torch.randn(1, 3, 3,
3)).size(),
torch.Size([1, 5, 3, 3]))
def test_chosen_inputs(self):
class Net(nn.Module):
def __init__(self, reduction):
super().__init__()
self.conv1 = nn.Conv2d(3, 3, kernel_size=1)
self.conv2 = nn.Conv2d(3, 3, kernel_size=1)
self.input = nn.InputChoice(2, n_chosen=2, reduction=reduction)
def forward(self, x):
x1 = self.conv1(x)
x2 = self.conv2(x)
return self.input([x1, x2])
for reduction in ['none', 'sum', 'mean', 'concat']:
model = self._convert_to_ir(Net(reduction))
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutator = mutators[0].bind_sampler(EnuemrateSampler())
model = mutator.apply(model)
result = self._get_converted_pytorch_model(model)(torch.randn(1, 3, 3, 3))
if reduction == 'none':
self.assertEqual(len(result), 2)
self.assertEqual(result[0].size(), torch.Size([1, 3, 3, 3]))
self.assertEqual(result[1].size(), torch.Size([1, 3, 3, 3]))
elif reduction == 'concat':
self.assertEqual(result.size(), torch.Size([1, 6, 3, 3]))
else:
self.assertEqual(result.size(), torch.Size([1, 3, 3, 3]))
def test_valuechoice_access_functional_expression(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.dropout_rate = nn.ValueChoice([[
1.05,
], [
1.1,
]])
def forward(self, x):
# if expression failed, the exception would be:
# ValueError: dropout probability has to be between 0 and 1, but got 1.05
return F.dropout(x, self.dropout_rate()[0] - .1)
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutator = mutators[0].bind_sampler(EnumerateSampler())
model1 = mutator.apply(model)
model2 = mutator.apply(model)
self._get_converted_pytorch_model(model1)(torch.randn(1, 3, 3, 3))
self.assertEqual(
self._get_converted_pytorch_model(model1)(torch.randn(1, 3, 3,
3)).size(),
torch.Size([1, 3, 3, 3]))
self.assertAlmostEqual(
self._get_converted_pytorch_model(model2)(torch.randn(
1, 3, 3, 3)).abs().sum().item(), 0)
def test_layer_choice(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.module = nn.LayerChoice([
nn.Conv2d(3, 3, kernel_size=1),
nn.Conv2d(3, 5, kernel_size=1)
])
def forward(self, x):
return self.module(x)
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutator = mutators[0].bind_sampler(EnumerateSampler())
model1 = mutator.apply(model)
model2 = mutator.apply(model)
self.assertEqual(
self._get_converted_pytorch_model(model1)(torch.randn(1, 3, 3,
3)).size(),
torch.Size([1, 3, 3, 3]))
self.assertEqual(
self._get_converted_pytorch_model(model2)(torch.randn(1, 3, 3,
3)).size(),
torch.Size([1, 5, 3, 3]))
def test_value_choice_as_parameter_shared(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, nn.ValueChoice([6, 8],
label='shared'), 1)
self.conv2 = nn.Conv2d(3, nn.ValueChoice([6, 8],
label='shared'), 1)
def forward(self, x):
return self.conv1(x) + self.conv2(x)
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutator = mutators[0].bind_sampler(EnumerateSampler())
model1 = mutator.apply(model)
model2 = mutator.apply(model)
self.assertEqual(
self._get_converted_pytorch_model(model1)(torch.randn(1, 3, 5,
5)).size(),
torch.Size([1, 6, 5, 5]))
self.assertEqual(
self._get_converted_pytorch_model(model2)(torch.randn(1, 3, 5,
5)).size(),
torch.Size([1, 8, 5, 5]))
def test_valuechoice_access_functional(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.dropout_rate = nn.ValueChoice([[
0.,
], [
1.,
]])
def forward(self, x):
return F.dropout(x, self.dropout_rate()[0])
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutator = mutators[0].bind_sampler(EnumerateSampler())
model1 = mutator.apply(model)
model2 = mutator.apply(model)
self._get_converted_pytorch_model(model1)(torch.randn(1, 3, 3, 3))
self.assertEqual(
self._get_converted_pytorch_model(model1)(torch.randn(1, 3, 3,
3)).size(),
torch.Size([1, 3, 3, 3]))
self.assertAlmostEqual(
self._get_converted_pytorch_model(model2)(torch.randn(
1, 3, 3, 3)).abs().sum().item(), 0)
def test_value_choice_as_parameter(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3,
nn.ValueChoice([6, 8]),
kernel_size=nn.ValueChoice([3, 5]))
def forward(self, x):
return self.conv(x)
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 2)
mutators[0].bind_sampler(EnumerateSampler())
mutators[1].bind_sampler(EnumerateSampler())
input = torch.randn(1, 3, 5, 5)
self.assertEqual(
self._get_converted_pytorch_model(mutators[1].apply(
mutators[0].apply(model)))(input).size(),
torch.Size([1, 6, 3, 3]))
self.assertEqual(
self._get_converted_pytorch_model(mutators[1].apply(
mutators[0].apply(model)))(input).size(),
torch.Size([1, 8, 1, 1]))
def _get_model_and_mutators():
base_model = Net()
script_module = torch.jit.script(base_model)
base_model_ir = convert_to_graph(script_module, base_model)
base_model_ir.training_config = DebugTraining()
mutators = process_inline_mutation(base_model_ir)
return base_model_ir, mutators
def test_shared(self):
class Net(nn.Module):
def __init__(self, shared=True):
super().__init__()
labels = ['x', 'x'] if shared else [None, None]
self.module1 = nn.LayerChoice([
nn.Conv2d(3, 3, kernel_size=1),
nn.Conv2d(3, 5, kernel_size=1)
],
label=labels[0])
self.module2 = nn.LayerChoice([
nn.Conv2d(3, 3, kernel_size=1),
nn.Conv2d(3, 5, kernel_size=1)
],
label=labels[1])
def forward(self, x):
return self.module1(x) + self.module2(x)
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
sampler = RandomSampler()
mutator = mutators[0].bind_sampler(sampler)
self.assertEqual(
self._get_converted_pytorch_model(mutator.apply(model))(
torch.randn(1, 3, 3, 3)).size(0), 1)
self.assertEqual(sampler.counter, 1)
model = self._convert_to_ir(Net(shared=False))
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 2)
sampler = RandomSampler()
# repeat test. Expectation: sometimes succeeds, sometimes fails.
failed_count = 0
for i in range(30):
model_new = model
for mutator in mutators:
model_new = mutator.bind_sampler(sampler).apply(model_new)
self.assertEqual(sampler.counter, 2 * (i + 1))
try:
self._get_converted_pytorch_model(model_new)(torch.randn(
1, 3, 3, 3))
except RuntimeError:
failed_count += 1
self.assertGreater(failed_count, 0)
self.assertLess(failed_count, 30)
def test_valuechoice_access(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
vc = nn.ValueChoice([(6, 3), (8, 5)])
self.conv = nn.Conv2d(3, vc[0], kernel_size=vc[1])
def forward(self, x):
return self.conv(x)
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutators[0].bind_sampler(EnumerateSampler())
input = torch.randn(1, 3, 5, 5)
self.assertEqual(
self._get_converted_pytorch_model(
mutators[0].apply(model))(input).size(),
torch.Size([1, 6, 3, 3]))
self.assertEqual(
self._get_converted_pytorch_model(
mutators[0].apply(model))(input).size(),
torch.Size([1, 8, 1, 1]))
class Net2(nn.Module):
def __init__(self):
super().__init__()
choices = [{'b': [3], 'bp': [6]}, {'b': [6], 'bp': [12]}]
self.conv = nn.Conv2d(
3,
nn.ValueChoice(choices, label='a')['b'][0], 1)
self.conv1 = nn.Conv2d(
nn.ValueChoice(choices, label='a')['bp'][0], 3, 1)
def forward(self, x):
x = self.conv(x)
return self.conv1(torch.cat((x, x), 1))
model = self._convert_to_ir(Net2())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutators[0].bind_sampler(EnumerateSampler())
input = torch.randn(1, 3, 5, 5)
self._get_converted_pytorch_model(mutators[0].apply(model))(input)
def test_value_choice(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.index = nn.ValueChoice([0, 1])
self.conv = MutableConv()
def forward(self, x):
return self.conv(x, self.index())
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 1)
mutator = mutators[0].bind_sampler(EnuemrateSampler())
model1 = mutator.apply(model)
model2 = mutator.apply(model)
self.assertEqual(self._get_converted_pytorch_model(model1)(torch.randn(1, 3, 3, 3)).size(),
torch.Size([1, 3, 3, 3]))
self.assertEqual(self._get_converted_pytorch_model(model2)(torch.randn(1, 3, 3, 3)).size(),
torch.Size([1, 5, 3, 3]))
def test_value_choice_in_layer_choice(self):
class Net(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.LayerChoice([
nn.Linear(3, nn.ValueChoice([10, 20])),
nn.Linear(3, nn.ValueChoice([30, 40]))
])
def forward(self, x):
return self.linear(x)
model = self._convert_to_ir(Net())
mutators = process_inline_mutation(model)
self.assertEqual(len(mutators), 3)
sz_counter = Counter()
sampler = RandomSampler()
for i in range(100):
model_new = model
for mutator in mutators:
model_new = mutator.bind_sampler(sampler).apply(model_new)
sz_counter[self._get_converted_pytorch_model(model_new)(
torch.randn(1, 3)).size(1)] += 1
self.assertEqual(len(sz_counter), 4)
def _get_model_with_mutators(self, pytorch_model):
model = self._convert_to_ir(pytorch_model)
mutators = process_inline_mutation(model)
return model, mutators
