#!/usr/bin/env python3 # Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang) # # See ../../../../LICENSE for clarification regarding multiple authors # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ To run this file, do: cd icefall/egs/librispeech/ASR python ./lstm_transducer_stateless/test_scaling_converter.py """ import copy import torch from scaling import ( ScaledConv1d, ScaledConv2d, ScaledEmbedding, ScaledLinear, ScaledLSTM, ) from scaling_converter import ( convert_scaled_to_non_scaled, scaled_conv1d_to_conv1d, scaled_conv2d_to_conv2d, scaled_embedding_to_embedding, scaled_linear_to_linear, scaled_lstm_to_lstm, ) from train import get_params, get_transducer_model def get_model(): params = get_params() params.vocab_size = 500 params.blank_id = 0 params.context_size = 2 params.unk_id = 2 params.encoder_dim = 512 params.rnn_hidden_size = 1024 params.num_encoder_layers = 12 params.aux_layer_period = -1 model = get_transducer_model(params) return model def test_scaled_linear_to_linear(): N = 5 in_features = 10 out_features = 20 for bias in [True, False]: scaled_linear = ScaledLinear( in_features=in_features, out_features=out_features, bias=bias, ) linear = scaled_linear_to_linear(scaled_linear) x = torch.rand(N, in_features) y1 = scaled_linear(x) y2 = linear(x) assert torch.allclose(y1, y2) jit_scaled_linear = torch.jit.script(scaled_linear) jit_linear = torch.jit.script(linear) y3 = jit_scaled_linear(x) y4 = jit_linear(x) assert torch.allclose(y3, y4) assert torch.allclose(y1, y4) def test_scaled_conv1d_to_conv1d(): in_channels = 3 for bias in [True, False]: scaled_conv1d = ScaledConv1d( in_channels, 6, kernel_size=1, stride=1, padding=0, bias=bias, ) conv1d = scaled_conv1d_to_conv1d(scaled_conv1d) x = torch.rand(20, in_channels, 10) y1 = scaled_conv1d(x) y2 = conv1d(x) assert torch.allclose(y1, y2) jit_scaled_conv1d = torch.jit.script(scaled_conv1d) jit_conv1d = torch.jit.script(conv1d) y3 = jit_scaled_conv1d(x) y4 = jit_conv1d(x) assert torch.allclose(y3, y4) assert torch.allclose(y1, y4) def test_scaled_conv2d_to_conv2d(): in_channels = 1 for bias in [True, False]: scaled_conv2d = ScaledConv2d( in_channels=in_channels, out_channels=3, kernel_size=3, padding=1, bias=bias, ) conv2d = scaled_conv2d_to_conv2d(scaled_conv2d) x = torch.rand(20, in_channels, 10, 20) y1 = scaled_conv2d(x) y2 = conv2d(x) assert torch.allclose(y1, y2) jit_scaled_conv2d = torch.jit.script(scaled_conv2d) jit_conv2d = torch.jit.script(conv2d) y3 = jit_scaled_conv2d(x) y4 = jit_conv2d(x) assert torch.allclose(y3, y4) assert torch.allclose(y1, y4) def test_scaled_embedding_to_embedding(): scaled_embedding = ScaledEmbedding( num_embeddings=500, embedding_dim=10, padding_idx=0, ) embedding = scaled_embedding_to_embedding(scaled_embedding) for s in [10, 100, 300, 500, 800, 1000]: x = torch.randint(low=0, high=500, size=(s,)) scaled_y = scaled_embedding(x) y = embedding(x) assert torch.equal(scaled_y, y) def test_scaled_lstm_to_lstm(): input_size = 512 batch_size = 20 for bias in [True, False]: for hidden_size in [512, 1024]: scaled_lstm = ScaledLSTM( input_size=input_size, hidden_size=hidden_size, num_layers=1, bias=bias, proj_size=0 if hidden_size == input_size else input_size, ) lstm = scaled_lstm_to_lstm(scaled_lstm) x = torch.rand(200, batch_size, input_size) h0 = torch.randn(1, batch_size, input_size) c0 = torch.randn(1, batch_size, hidden_size) y1, (h1, c1) = scaled_lstm(x, (h0, c0)) y2, (h2, c2) = lstm(x, (h0, c0)) assert torch.allclose(y1, y2) assert torch.allclose(h1, h2) assert torch.allclose(c1, c2) jit_scaled_lstm = torch.jit.trace(lstm, (x, (h0, c0))) y3, (h3, c3) = jit_scaled_lstm(x, (h0, c0)) assert torch.allclose(y1, y3) assert torch.allclose(h1, h3) assert torch.allclose(c1, c3) def test_convert_scaled_to_non_scaled(): for inplace in [False, True]: model = get_model() model.eval() orig_model = copy.deepcopy(model) converted_model = convert_scaled_to_non_scaled(model, inplace=inplace) model = orig_model # test encoder N = 2 T = 100 vocab_size = model.decoder.vocab_size x = torch.randn(N, T, 80, dtype=torch.float32) x_lens = torch.full((N,), x.size(1)) e1, e1_lens, _ = model.encoder(x, x_lens) e2, e2_lens, _ = converted_model.encoder(x, x_lens) assert torch.all(torch.eq(e1_lens, e2_lens)) assert torch.allclose(e1, e2), (e1 - e2).abs().max() # test decoder U = 50 y = torch.randint(low=1, high=vocab_size - 1, size=(N, U)) d1 = model.decoder(y) d2 = model.decoder(y) assert torch.allclose(d1, d2) # test simple projection lm1 = model.simple_lm_proj(d1) am1 = model.simple_am_proj(e1) lm2 = converted_model.simple_lm_proj(d2) am2 = converted_model.simple_am_proj(e2) assert torch.allclose(lm1, lm2) assert torch.allclose(am1, am2) # test joiner e = torch.rand(2, 3, 4, 512) d = torch.rand(2, 3, 4, 512) j1 = model.joiner(e, d) j2 = converted_model.joiner(e, d) assert torch.allclose(j1, j2) @torch.no_grad() def main(): test_scaled_linear_to_linear() test_scaled_conv1d_to_conv1d() test_scaled_conv2d_to_conv2d() test_scaled_embedding_to_embedding() test_scaled_lstm_to_lstm() test_convert_scaled_to_non_scaled() if __name__ == "__main__": torch.manual_seed(20220730) main()