def test_amp_conversion_rnn(amp_tests):
with mx.Context(mx.gpu(0)):
model = nn.HybridSequential()
model.add(rnn.LSTM(hidden_size=10, num_layers=2, bidirectional=True))
model.add(nn.Dense(2))
model.initialize()
model.hybridize()
out = model(mx.nd.ones((2, 3, 4)))
new_model = amp.convert_hybrid_block(model)
out2 = new_model(mx.nd.ones((2, 3, 4)))
mx.test_utils.assert_almost_equal(out.asnumpy(), out2.asnumpy(), atol=1e-2, rtol=1e-2)
def check_amp_convert_hybrid_block():
# Test conversion for hybrid block on CPU
model_cpu = get_model("resnet50_v1")
model_cpu.collect_params().initialize(ctx=mx.cpu())
model_cpu.hybridize()
model_cpu(mx.nd.random.uniform(0, 1, shape=(1, 3, 224, 224), ctx=mx.cpu()))
converted_model_cpu = amp.convert_hybrid_block(model_cpu)
# Test with real world model, default inputs for convert_hybrid_block
model = get_model("resnet50_v1")
model.collect_params().initialize(ctx=mx.gpu())
model.hybridize()
model(mx.nd.zeros((1, 3, 224, 224)))
converted_model = amp.convert_hybrid_block(model)
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224),
dtype=np.float32))
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224),
dtype=np.float32))
# Test with real world model, tweak inputs for convert_hybrid_block
converted_model = amp.convert_hybrid_block(model, target_dtype="float16",
target_dtype_ops=["Convolution"])
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224),
dtype=np.float32))
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224),
dtype=np.float32))
# Check symbolic block
dir_path = os.path.dirname(os.path.realpath(__file__))
model_path = os.path.join(dir_path, 'model')
if not os.path.isdir(model_path):
os.mkdir(model_path)
prefix, epoch = download_model("imagenet1k-resnet-18", dst_dir=model_path)
net = SymbolBlock.imports(os.path.join(model_path, "imagenet1k-resnet-18-symbol.json"),
input_names=["data", "softmax_label"],
param_file=os.path.join(model_path, "imagenet1k-resnet-18-0000.params"))
net.collect_params().reset_ctx(ctx=mx.gpu())
net.hybridize()
net(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1,)))
converted_model = amp.convert_hybrid_block(net)
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1,)))
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1,)))
# Check symbolic block, tweaked inputs
converted_model = amp.convert_hybrid_block(net, target_dtype="float16", target_dtype_ops=["Convolution"])
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1, )))
result = converted_model.forward(mx.nd.zeros((1, 3, 224, 224)), mx.nd.zeros((1, )))
params = converted_model.collect_params()
assert params["stage2_unit1_conv2_weight"].dtype == np.float32
# Pass cast_optional_params as True to convert_hybrid_block
converted_model = amp.convert_hybrid_block(net, target_dtype="float16", target_dtype_ops=["Convolution"],
cast_optional_params=True)
params = converted_model.collect_params()
assert params["stage2_unit1_conv2_weight"].dtype == np.float16
