(self)
| 50 | self.assertEqual(outputs[2].shape, (input_shape[0],) + input_shape[2:]) # (1, H, W, [D]) |
| 51 | |
| 52 | def test_ema(self): |
| 53 | layer = EMAQuantizer(spatial_dims=2, num_embeddings=2, embedding_dim=2, epsilon=0, decay=0) |
| 54 | original_weight_0 = layer.embedding.weight[0].clone() |
| 55 | original_weight_1 = layer.embedding.weight[1].clone() |
| 56 | x_0 = original_weight_0 |
| 57 | x_0 = x_0.unsqueeze(0).unsqueeze(-1).unsqueeze(-1) |
| 58 | x_0 = x_0.repeat(1, 1, 1, 2) + 0.001 |
| 59 | |
| 60 | x_1 = original_weight_1 |
| 61 | x_1 = x_1.unsqueeze(0).unsqueeze(-1).unsqueeze(-1) |
| 62 | x_1 = x_1.repeat(1, 1, 1, 2) |
| 63 | |
| 64 | x = torch.cat([x_0, x_1], dim=0) |
| 65 | layer = layer.train() |
| 66 | _ = layer(x) |
| 67 | |
| 68 | self.assertTrue(all(layer.embedding.weight[0] != original_weight_0)) |
| 69 | self.assertTrue(all(layer.embedding.weight[1] == original_weight_1)) |
| 70 | |
| 71 | |
| 72 | class TestVectorQuantizer(unittest.TestCase): |
nothing calls this directly
no test coverage detected