def test_script(self):
net = AutoEncoder(spatial_dims=2,
in_channels=1,
out_channels=1,
channels=(4, 8),
strides=(2, 2))
test_data = torch.randn(2, 1, 32, 32)
test_script_save(net, test_data)
def test_script(self):
net = AutoEncoder(dimensions=2,
in_channels=1,
out_channels=1,
channels=(4, 8),
strides=(2, 2))
test_data = torch.randn(2, 1, 32, 32)
out_orig, out_reloaded = test_script_save(net, test_data)
assert torch.allclose(out_orig, out_reloaded)
def train(dict_key_for_training, max_epochs=10, learning_rate=1e-3):
model = AutoEncoder(
dimensions=2,
in_channels=1,
out_channels=1,
channels=(4, 8, 16, 32),
strides=(2, 2, 2, 2),
).to(device)
# print( dict_key_for_training )
# Create loss fn and optimiser
loss_function = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), learning_rate)
epoch_loss_values = []
t = trange(max_epochs,
desc=f"{dict_key_for_training} -- epoch 0, avg loss: inf",
leave=True)
for epoch in t:
model.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step += 1
inputs = batch_data[dict_key_for_training].to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs,
batch_data[dict_key_for_training].to(device))
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_loss /= step
print('OutOfLoop')
epoch_loss_values.append(epoch_loss)
t.set_description(f"{dict_key_for_training} -- epoch {epoch + 1}" +
f", average loss: {epoch_loss:.4f}")
return model, epoch_loss_values
def test_channel_stride_difference(self):
with self.assertRaises(ValueError):
net = AutoEncoder(**TEST_CASE_FAIL)
def test_shape(self, input_param, input_shape, expected_shape):
net = AutoEncoder(**input_param).to(device)
with eval_mode(net):
result = net.forward(torch.randn(input_shape).to(device))
self.assertEqual(result.shape, expected_shape)