| 27 | |
| 28 | def run_test(batch_size=64, train_steps=200, device="cuda:0"): |
| 29 | class _TestBatch(Dataset): |
| 30 | def __init__(self, transforms): |
| 31 | self.transforms = transforms |
| 32 | |
| 33 | def __getitem__(self, _unused_id): |
| 34 | im, seg = create_test_image_2d(128, 128, noise_max=1, num_objs=4, num_seg_classes=1) |
| 35 | seed = np.random.randint(2147483647) |
| 36 | self.transforms.set_random_state(seed=seed) |
| 37 | im = self.transforms(im) |
| 38 | self.transforms.set_random_state(seed=seed) |
| 39 | seg = self.transforms(seg) |
| 40 | return im, seg |
| 41 | |
| 42 | def __len__(self): |
| 43 | return train_steps |
| 44 | |
| 45 | net = UNet( |
| 46 | spatial_dims=2, in_channels=1, out_channels=1, channels=(4, 8, 16, 32), strides=(2, 2, 2), num_res_units=2 |