def test_unet_3d_decoder_creation(self, model_id):
"""Test creation of UNet 3D decoder."""
# Create test input for decoders based on input model_id.
input_specs = {}
for level in range(model_id):
input_specs[str(level + 1)] = tf.TensorShape([
1, 128 // (2**level), 128 // (2**level), 128 // (2**level), 1
])
network = decoders.UNet3DDecoder(model_id=model_id,
input_specs=input_specs,
use_sync_bn=True,
use_batch_normalization=True,
use_deconvolution=True)
model_config = semantic_segmentation_3d_exp.SemanticSegmentationModel3D(
)
model_config.num_classes = 2
model_config.num_channels = 1
model_config.input_size = [None, None, None]
model_config.decoder = decoders_cfg.Decoder(
type='unet_3d_decoder',
unet_3d_decoder=decoders_cfg.UNet3DDecoder(model_id=model_id))
factory_network = factory.build_decoder(input_specs=input_specs,
model_config=model_config)
network_config = network.get_config()
factory_network_config = factory_network.get_config()
print(network_config)
print(factory_network_config)
self.assertEqual(network_config, factory_network_config)
def seg_unet3d_test() -> cfg.ExperimentConfig:
"""Image segmentation on a dummy dataset with 3D UNet for testing purpose."""
train_batch_size = 2
eval_batch_size = 2
steps_per_epoch = 10
config = cfg.ExperimentConfig(
task=SemanticSegmentation3DTask(
model=SemanticSegmentationModel3D(
num_classes=2,
input_size=[32, 32, 32],
num_channels=2,
backbone=backbones.Backbone(
type='unet_3d', unet_3d=backbones.UNet3D(model_id=2)),
decoder=decoders.Decoder(
type='unet_3d_decoder',
unet_3d_decoder=decoders.UNet3DDecoder(model_id=2)),
head=SegmentationHead3D(num_convs=0, num_classes=2),
norm_activation=common.NormActivation(
activation='relu', use_sync_bn=False)),
train_data=DataConfig(
input_path='train.tfrecord',
num_classes=2,
input_size=[32, 32, 32],
num_channels=2,
is_training=True,
global_batch_size=train_batch_size),
validation_data=DataConfig(
input_path='val.tfrecord',
num_classes=2,
input_size=[32, 32, 32],
num_channels=2,
is_training=False,
global_batch_size=eval_batch_size),
losses=Losses(loss_type='adaptive')),
trainer=cfg.TrainerConfig(
steps_per_loop=steps_per_epoch,
summary_interval=steps_per_epoch,
checkpoint_interval=steps_per_epoch,
train_steps=10,
validation_steps=10,
validation_interval=steps_per_epoch,
optimizer_config=optimization.OptimizationConfig({
'optimizer': {
'type': 'sgd',
},
'learning_rate': {
'type': 'constant',
'constant': {
'learning_rate': 0.000001
}
}
})),
restrictions=[
'task.train_data.is_training != None',
'task.validation_data.is_training != None'
])
return config
class SemanticSegmentationModel3D(hyperparams.Config):
"""Semantic segmentation model config."""
num_classes: int = 0
num_channels: int = 1
input_size: List[int] = dataclasses.field(default_factory=list)
min_level: int = 3
max_level: int = 6
head: SegmentationHead3D = SegmentationHead3D()
backbone: backbones.Backbone = backbones.Backbone(
type='unet_3d', unet_3d=backbones.UNet3D())
decoder: decoders.Decoder = decoders.Decoder(
type='unet_3d_decoder', unet_3d_decoder=decoders.UNet3DDecoder())
norm_activation: common.NormActivation = common.NormActivation()