def find_2d_configuration():
# lets compute a reference for 3D
# we select hyperparameters here so that we get approximately the same patch size as we would get with the
# regular unet. This is just my choice. You can do whatever you want
# These default hyperparemeters will then be used by the experiment planner
# since this is more parameter intensive than the UNet, we will test a configuration that has a lot of parameters
# herefore we copy the UNet configuration for Task003_Liver
cudnn.deterministic = False
cudnn.benchmark = True
patch_size = (512, 512)
max_num_features = 512
num_modalities = 1
num_classes = 3
batch_size = 12
# now we fiddle with the network specific hyperparameters until everything just barely fits into a titanx
blocks_per_stage_encoder = FabiansUNet.default_blocks_per_stage_encoder
blocks_per_stage_decoder = FabiansUNet.default_blocks_per_stage_decoder
initial_num_features = 30
# we neeed to add a [1, 1, 1] for the res unet because in this implementation all stages of the encoder can have a stride
pool_op_kernel_sizes = [[1, 1],
[2, 2],
[2, 2],
[2, 2],
[2, 2],
[2, 2],
[2, 2],
[2, 2]]
conv_op_kernel_sizes = [[3, 3],
[3, 3],
[3, 3],
[3, 3],
[3, 3],
[3, 3],
[3, 3],
[3, 3]]
unet = FabiansUNet(num_modalities, initial_num_features, blocks_per_stage_encoder[:len(conv_op_kernel_sizes)], 2,
pool_op_kernel_sizes, conv_op_kernel_sizes,
get_default_network_config(2, dropout_p=None), num_classes,
blocks_per_stage_decoder[:len(conv_op_kernel_sizes)-1], False, False,
max_features=max_num_features).cuda()
optimizer = SGD(unet.parameters(), lr=0.1, momentum=0.95)
loss = DC_and_CE_loss({'batch_dice': True, 'smooth': 1e-5, 'do_bg': False}, {})
dummy_input = torch.rand((batch_size, num_modalities, *patch_size)).cuda()
dummy_gt = (torch.rand((batch_size, 1, *patch_size)) * num_classes).round().clamp_(0, 2).cuda().long()
for _ in range(20):
optimizer.zero_grad()
skips = unet.encoder(dummy_input)
print([i.shape for i in skips])
output = unet.decoder(skips)
l = loss(output, dummy_gt)
l.backward()
optimizer.step()
if _ == 0:
torch.cuda.empty_cache()
# that should do. Now take the network hyperparameters and insert them in FabiansUNet.compute_approx_vram_consumption
# whatever number this spits out, save it to FabiansUNet.use_this_for_batch_size_computation_2D
print(FabiansUNet.compute_approx_vram_consumption(patch_size, initial_num_features, max_num_features, num_modalities,
num_classes, pool_op_kernel_sizes,
blocks_per_stage_encoder[:len(conv_op_kernel_sizes)],
blocks_per_stage_decoder[:len(conv_op_kernel_sizes)-1], 2, batch_size))
def find_2d_configuration():
cudnn.benchmark = True
cudnn.deterministic = False
conv_op_kernel_sizes = ((3, 3), (3, 3), (3, 3), (3, 3), (3, 3), (3, 3),
(3, 3))
pool_op_kernel_sizes = ((1, 1), (2, 2), (2, 2), (2, 2), (2, 2), (2, 2),
(2, 2))
patch_size = (256, 256)
base_num_features = 32
input_modalities = 4
blocks_per_stage_encoder = (1, 3, 4, 6, 6, 6, 6)
blocks_per_stage_decoder = (2, 2, 2, 2, 2, 2)
feat_map_mult_on_downscale = 2
num_classes = 5
max_features = 512
batch_size = 50
unet = FabiansPreActUNet(input_modalities,
base_num_features,
blocks_per_stage_encoder,
feat_map_mult_on_downscale,
pool_op_kernel_sizes,
conv_op_kernel_sizes,
get_default_network_config(2, dropout_p=None),
num_classes,
blocks_per_stage_decoder,
True,
False,
max_features=max_features).cuda()
scaler = GradScaler()
optimizer = SGD(unet.parameters(), lr=0.1, momentum=0.95)
print(
unet.compute_approx_vram_consumption(
patch_size, base_num_features, max_features, input_modalities,
num_classes, pool_op_kernel_sizes, blocks_per_stage_encoder,
blocks_per_stage_decoder, feat_map_mult_on_downscale, batch_size))
loss = DC_and_CE_loss({
'batch_dice': True,
'smooth': 1e-5,
'do_bg': False
}, {})
dummy_input = torch.rand(
(batch_size, input_modalities, *patch_size)).cuda()
dummy_gt = (torch.rand(
(batch_size, 1, *patch_size)) * num_classes).round().clamp_(
0, num_classes - 1).cuda().long()
for i in range(10):
optimizer.zero_grad()
with autocast():
skips = unet.encoder(dummy_input)
print([i.shape for i in skips])
output = unet.decoder(skips)[0]
l = loss(output, dummy_gt)
print(l.item())
scaler.scale(l).backward()
scaler.step(optimizer)
scaler.update()
with autocast():
import hiddenlayer as hl
g = hl.build_graph(unet, dummy_input, transforms=None)
g.save("/home/fabian/test_arch.pdf")