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)
