def test_weight_decay(module, name, steps: bool): tl.build(module, torch.randn(1, 20)) for _ in range(steps): output = module(torch.randn(5, 20)) output.sum().backward() check_gradient_correctness(module, name) module.zero_grad() output = module(torch.rand(5, 20)) output.sum().backward() check_gradient_correctness(module, name)
def test_text_cnn(): model = torch.nn.Sequential( torchlayers.Conv(64), # specify ONLY out_channels torch.nn.ReLU(), # use torch.nn wherever you wish torchlayers.BatchNorm(), # BatchNormNd inferred from input torchlayers.Conv(128), # Default kernel_size equal to 3 torchlayers.ReLU(), torchlayers.Conv(256, kernel_size=11), # "same" padding as default torchlayers.GlobalMaxPool(), # Known from Keras torchlayers.Linear(10), # Output for 10 classes ) torchlayers.build(model, torch.randn(2, 300, 1))
def test_gru(): model = tl.build( tl.GRU(hidden_size=40, dropout=0.5, batch_first=True, num_layers=3), torch.randn(5, 3, 10), ) output, _ = model(torch.randn(5, 3, 10)) assert output.shape == (5, 3, 40)
def test_basic_jit(): inputs = torch.randn(16, 3, 32, 32) layer = tl.build(tl.Conv(64), inputs) output = layer(inputs) new_model = torch.jit.script(layer) new_output = new_model(inputs) assert torch.allclose(output, new_output)
def test_save(): inputs = torch.randn(16, 32) temp = pathlib.Path(tempfile.gettempdir()) layer = tl.build(tl.Linear(64), inputs) output = layer(inputs) torch.save(layer, temp / "linear_model.pt") new_layer = torch.load(temp / "linear_model.pt") new_output = new_layer(inputs) assert torch.allclose(output, new_output)
def test_convolution_save(): inputs = torch.randn(16, 3, 32, 32) temp = pathlib.Path(tempfile.gettempdir()) layer = tl.build(tl.Conv2d(64, kernel_size=3), inputs) output = layer(inputs) torch.save(layer, temp / "conv_model.pt") new_layer = torch.load(temp / "conv_model.pt") new_output = new_layer(inputs) assert torch.allclose(output, new_output)
def test_classification(classification_model): classification_model = torchlayers.build(classification_model, torch.randn(16, 3, 28, 28)) optimizer = torch.optim.Adam(classification_model.parameters()) criterion = torch.nn.CrossEntropyLoss() for _ in range(16): output = classification_model(torch.randn(16, 3, 28, 28)) loss = criterion(output, torch.randint(10, (16, ))) loss.backward() optimizer.zero_grad()
def test_functionality(model): # Initialize model = tl.build(model, torch.randn(16, 3, 28, 28)) optimizer = torch.optim.Adam(model.parameters()) criterion = torch.nn.CrossEntropyLoss() for _ in range(16): output = model(torch.randn(16, 3, 28, 28)) loss = criterion(output, torch.randint(2, (16, ))) loss.backward() optimizer.zero_grad()
def test_basic_jit_save(): inputs = torch.randn(16, 3, 32, 32) temp = pathlib.Path(tempfile.gettempdir()) layer = tl.build(tl.Conv(64), inputs) output = layer(inputs) new_model = torch.jit.script(layer) torch.jit.save(new_model, str(temp / "jit.pt")) loaded_model = torch.jit.load(str(temp / "jit.pt")) new_output = loaded_model(inputs) assert torch.allclose(output, new_output)
def test_autoencoder(autoencoder_model): autoencoder_model = torchlayers.build(autoencoder_model, torch.randn(1, 3, 256, 256)) optimizer = torch.optim.Adam(autoencoder_model.parameters()) criterion = torch.nn.MSELoss() for _ in range(16): inputs = torch.randn(1, 3, 256, 256) output = autoencoder_model(inputs) loss = criterion(output, inputs) loss.backward() optimizer.zero_grad()
def test_same_padding_1d(): inputs = torch.randn(1, 5, 125) for kernel_size, stride, dilation in itertools.product( range(1, 20, 2), *[range(1, 20, 2) for _ in range(2)]): classification_model = torchlayers.build( torchlayers.Conv( 5, kernel_size=kernel_size, stride=stride, dilation=dilation, padding="same", ), inputs, ) output = classification_model(inputs) assert output.shape == inputs.shape
def test_same_padding_2d(): inputs = torch.randn(1, 5, 100, 100) for kernel_size, stride, dilation in itertools.product( range(1, 15, 2), *[range(1, 8, 2) for _ in range(2)]): for kernel_size2, stride2, dilation2 in itertools.product( range(1, 8, 2), *[range(1, 8, 2) for _ in range(2)]): classification_model = torchlayers.build( torchlayers.Conv( 5, kernel_size=(kernel_size, kernel_size2), stride=(stride, stride2), dilation=(dilation, dilation2), padding="same", ), inputs, ) output = classification_model(inputs) assert output.shape == inputs.shape
def test_residual_bottleneck(): model = torchlayers.build(torchlayers.InvertedResidualBottleneck(), torch.randn(1, 128, 128, 128))
def test_conv_pixel_shuffle(): model = torchlayers.build( torchlayers.ConvPixelShuffle(out_channels=512, upscale_factor=2), torch.randn(1, 256, 128, 128), )
def test_functionality(model): # Initialize model = torchlayers.build(model, torch.randn(16, 3, 28, 28)) optimizer = torch.optim.Adam(model.parameters()) criterion = torch.nn.CrossEntropyLoss()
def test_last_dimension_linear(): module = tl.Linear(64) module = tl.build(module, torch.randn(1, 3, 32, 32)) assert module(torch.randn(2, 6, 24, 32)).shape == (2, 6, 24, 64)
def test_lstm(): model = torchlayers.build(torchlayers.LSTM(hidden_size=20), torch.randn(5, 3, 10)) output, (_, _) = model(torch.randn(5, 3, 10)) assert output.shape == (5, 3, 20)
def test_custom_inferable_build(): layer = CustomLinear(32) layer = tl.build(layer, torch.rand(16, 64)) assert layer.some_params.shape == (2, 32)