def initialize_network(self):
"""
- momentum 0.99
- SGD instead of Adam
- self.lr_scheduler = None because we do poly_lr
- deep supervision = True
- i am sure I forgot something here
Known issue: forgot to set neg_slope=0 in InitWeights_He; should not make a difference though
:return:
"""
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
if self.net == '3d_nas':
self.network = Nas_UNet(
4, 4, self.num_input_channels,
self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes), self.conv_per_stage, 2,
conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs, net_nonlin,
net_nonlin_kwargs, True, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
self.architect = Architect(self.network, self.loss)
elif self.net == '3d_nas_2C':
self.network = Nas_UNet(
2, 2, self.num_input_channels,
self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes), self.conv_per_stage, 2,
conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs, net_nonlin,
net_nonlin_kwargs, True, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
self.architect = Architect(self.network, self.loss)
else:
self.network = Generic_UNet(
self.num_input_channels,
self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes), self.conv_per_stage, 2,
conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs, net_nonlin,
net_nonlin_kwargs, True, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
if torch.cuda.is_available():
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def initialize_network(self):
self.base_num_features = 24 # otherwise we run out of VRAM
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
self.network = Generic_UNet(
self.num_input_channels, self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes), 3, 2, conv_op, norm_op,
norm_op_kwargs, dropout_op, dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, True, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def initialize_network(self):
"""
changed deep supervision to False
:return:
"""
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
self.network = Generic_UNet(
self.num_input_channels, self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes), self.conv_per_stage, 2,
conv_op, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, False, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
if torch.cuda.is_available():
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def initialize_network(self):
"""
This is specific to the U-Net and must be adapted for other network architectures
:return:
"""
# self.print_to_log_file(self.net_num_pool_op_kernel_sizes)
# self.print_to_log_file(self.net_conv_kernel_sizes)
net_numpool = len(self.net_num_pool_op_kernel_sizes)
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
self.network = Generic_UNet(
self.num_input_channels, self.base_num_features, self.num_classes,
net_numpool, self.conv_per_stage, 2, conv_op, norm_op,
norm_op_kwargs, dropout_op, dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, False, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
self.network.inference_apply_nonlin = softmax_helper
if torch.cuda.is_available():
self.network.cuda()
def initialize_network(self):
"""
replace genericUNet with the implementation of above for super speeds
"""
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
if self.net=='3d_nas':
self.network=Nas_UNet(self.num_input_channels, self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes),
self.conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, True, False, lambda x: x, InitWeights_He(1e-2),
self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True, True)
self.architect = Architect(self.network, self.loss)
else:
self.network = Generic_UNet_DP(self.num_input_channels, self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes),
self.conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, True, False, InitWeights_He(1e-2),
self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True, True)
if torch.cuda.is_available():
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def initialize_network(self):
if self.threeD:
cfg = get_default_network_config(3, None, norm_type="in")
else:
cfg = get_default_network_config(1, None, norm_type="in")
stage_plans = self.plans['plans_per_stage'][self.stage]
conv_kernel_sizes = stage_plans['conv_kernel_sizes']
blocks_per_stage_encoder = stage_plans['num_blocks_encoder']
blocks_per_stage_decoder = stage_plans['num_blocks_decoder']
pool_op_kernel_sizes = stage_plans['pool_op_kernel_sizes']
self.network = FabiansUNet(self.num_input_channels,
self.base_num_features,
blocks_per_stage_encoder, 2,
pool_op_kernel_sizes, conv_kernel_sizes,
cfg, self.num_classes,
blocks_per_stage_decoder, True, False, 320,
InitWeights_He(1e-2))
self.network.cuda()
#vivian added
print(self.num_input_channels, self.base_num_features,
self.num_classes)
summary(self.network, (self.num_input_channels, 128, 128, 128))
self.network.inference_apply_nonlin = softmax_helper
def initialize_network(self):
"""
changed deep supervision to False
:return:
"""
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = FRN3D
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
raise NotImplementedError
norm_op = nn.BatchNorm2d
norm_op_kwargs = {'eps': 1e-6}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = Identity
net_nonlin_kwargs = {}
self.network = Generic_UNet(
self.num_input_channels, self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes), self.conv_per_stage, 2,
conv_op, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, True, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
if torch.cuda.is_available():
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def initialize_network_optimizer_and_scheduler(self):
"""
This is specific to the U-Net and must be adapted for other network architectures
:return:
"""
# self.print_to_log_file(self.net_num_pool_op_kernel_sizes)
# self.print_to_log_file(self.net_conv_kernel_sizes)
net_numpool = len(self.net_num_pool_op_kernel_sizes)
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
self.network = Generic_UNet(
self.num_input_channels, self.base_num_features, self.num_classes,
net_numpool, 2, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, False, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
self.optimizer = torch.optim.Adam(self.network.parameters(),
self.initial_lr,
weight_decay=self.weight_decay,
amsgrad=True)
self.lr_scheduler = lr_scheduler.ReduceLROnPlateau(
self.optimizer,
mode='min',
factor=0.2,
patience=self.lr_scheduler_patience,
verbose=True,
threshold=self.lr_scheduler_eps,
threshold_mode="abs")
self.network.cuda()
if torch.cuda.device_count() > 1:
torch.distributed.init_process_group(
backend='nccl',
init_method='tcp://localhost:23456',
rank=0,
world_size=1)
self.network = torch.nn.parallel.DistributedDataParallel(
self.network)
# self.network, self.optimizer = amp.initialize(self.network, self.optimizer, opt_level="O1")
# self.network = DistributedDataParallel(self.network)
print("Let's use", torch.cuda.device_count(), "GPUs!")
self.network.inference_apply_nonlin = softmax_helper
def initialize_network(self):
"""
- momentum 0.99
- SGD instead of Adam
- self.lr_scheduler = None because we do poly_lr
- deep supervision = True
- i am sure I forgot something here
Known issue: forgot to set neg_slope=0 in InitWeights_He; should not make a difference though
:return:
"""
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.BatchNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
#model = ResAxialAttentionUNet(AxialBlock, [1, 2, 4, 1], s=0.125, **kwargs)
"""self.network = ResAxialAttentionUNet([1, 1, 1, 1], self.num_input_channels,
self.num_classes, len(self.net_num_pool_op_kernel_sizes),
self.conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, False, lambda x: x, InitWeights_He(1e-2),
self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True,
True, s=0.0625, img_size=self.patch_size)"""
self.network = medt_net([1, 2, 4, 1], self.num_input_channels,
self.num_classes, len(self.net_num_pool_op_kernel_sizes),
self.conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, False, lambda x: x, InitWeights_He(1e-2),
self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True,
True, s=0.125, img_size=self.patch_size)
"""self.network = VNet_recons(self.num_input_channels, self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes),
self.conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, True, False, lambda x: x, InitWeights_He(1e-2),
self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True, True)
"""
if torch.cuda.is_available():
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
"""self.network = VNet_recons()
def initialize_network(self):
"""
- momentum 0.99
- SGD instead of Adam
- self.lr_scheduler = None because we do poly_lr
- deep supervision = True
- i am sure I forgot something here
Known issue: forgot to set neg_slope=0 in InitWeights_He; should not make a difference though
:return:
"""
assert self.threeD is False
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
print("sawNetTrainer.py:58", self.net_num_pool_op_kernel_sizes)
self.network = SAWNet(
self.num_input_channels,
self.base_num_features,
self.num_classes,
len(self.net_num_pool_op_kernel_sizes),
self.conv_per_stage,
2,
conv_op,
norm_op,
norm_op_kwargs,
dropout_op,
dropout_op_kwargs, # ds below
net_nonlin,
net_nonlin_kwargs,
False,
False,
softmax_helper,
InitWeights_He(1e-2),
self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes,
False,
True,
True)
if torch.cuda.is_available():
self.network.cuda()
def initialize_network(self):
"""
This is specific to the U-Net and must be adapted for other network architectures
:return:
"""
net_numpool = len(self.net_num_pool_op_kernel_sizes)
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
unet_kw = {
'input_channels': self.num_input_channels,
'base_num_features': self.base_num_features,
'num_classes': self.num_classes,
'num_pool': net_numpool,
'num_conv_per_stage': self.conv_per_stage,
'feat_map_mul_on_downscale': 2,
'conv_op': conv_op,
'norm_op': norm_op,
'norm_op_kwargs': norm_op_kwargs,
'dropout_op': dropout_op,
'dropout_op_kwargs': dropout_op_kwargs,
'nonlin': net_nonlin,
'nonlin_kwargs': net_nonlin_kwargs,
'deep_supervision': True,
'dropout_in_localization': False,
'final_nonlin': lambda x: x,
'weightInitializer': InitWeights_He(1e-2),
'pool_op_kernel_sizes': self.net_num_pool_op_kernel_sizes,
'conv_kernel_sizes': self.net_conv_kernel_sizes,
'upscale_logits': False,
'convolutional_pooling': True,
'convolutional_upsampling': True,
'max_num_features': None,
'seg_output_use_bias': False
}
self.network = CountingUNet(**unet_kw)
self.network.inference_apply_nonlin = softmax_helper
if torch.cuda.is_available():
self.network.cuda()
def initialize_network_optimizer_and_scheduler(self):
"""
This is specific to the U-Net and must be adapted for other network architectures
:return:
"""
# self.print_to_log_file(self.net_num_pool_op_kernel_sizes)
# self.print_to_log_file(self.net_conv_kernel_sizes)
net_numpool = len(self.net_num_pool_op_kernel_sizes)
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
self.network = Generic_UNet(
self.num_input_channels, self.base_num_features, self.num_classes,
net_numpool, 2, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, False, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
self.optimizer = torch.optim.Adam(self.network.parameters(),
self.initial_lr,
weight_decay=self.weight_decay,
amsgrad=True)
self.lr_scheduler = lr_scheduler.ReduceLROnPlateau(
self.optimizer,
mode='min',
factor=0.2,
patience=self.lr_scheduler_patience,
verbose=True,
threshold=self.lr_scheduler_eps,
threshold_mode="abs")
gpu_ids = [int(x) for x in self.gpu.split(",")]
if len(gpu_ids) > 1:
self.network = torch.nn.DataParallel(self.network.cuda(),
device_ids=gpu_ids)
else:
self.network = self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def __init__(self,
input_channels,
base_num_features,
num_classes,
num_pool,
num_conv_per_stage=2,
feat_map_mul_on_downscale=2,
conv_op=nn.Conv2d,
norm_op=nn.BatchNorm2d,
norm_op_kwargs=None,
dropout_op=nn.Dropout2d,
dropout_op_kwargs=None,
nonlin=nn.LeakyReLU,
nonlin_kwargs=None,
deep_supervision=True,
dropout_in_localization=False,
final_nonlin=softmax_helper,
weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=None,
conv_kernel_sizes=None,
upscale_logits=False,
convolutional_pooling=False,
convolutional_upsampling=False,
max_num_features=None,
basic_block=ConvDropoutNormNonlin,
seg_output_use_bias=False):
super(SAUNet, self).__init__(
input_channels, base_num_features, num_classes, num_pool,
num_conv_per_stage, feat_map_mul_on_downscale, conv_op, norm_op,
norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin,
nonlin_kwargs, deep_supervision, dropout_in_localization,
final_nonlin, weightInitializer, pool_op_kernel_sizes,
conv_kernel_sizes, upscale_logits, convolutional_pooling,
convolutional_upsampling, max_num_features, basic_block,
seg_output_use_bias)
output_features = base_num_features
for d in range(num_pool):
input_features = output_features
output_features = int(
np.round(output_features * feat_map_mul_on_downscale))
output_features = min(output_features, self.max_num_features)
print("{}: {}".format(d, input_features))
self.sau = SAUnit(output_features)
def __init__(self,
input_channels,
base_num_features,
num_classes,
num_pool,
num_conv_per_stage=2,
feat_map_mul_on_downscale=2,
conv_op=nn.Conv2d,
norm_op=nn.BatchNorm2d,
norm_op_kwargs=None,
dropout_op=nn.Dropout2d,
dropout_op_kwargs=None,
nonlin=nn.LeakyReLU,
nonlin_kwargs=None,
deep_supervision=True,
dropout_in_localization=False,
weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=None,
conv_kernel_sizes=None,
upscale_logits=False,
convolutional_pooling=False,
convolutional_upsampling=False,
max_num_features=None):
"""
As opposed to the Generic_UNet, this class will compute parts of the loss function in the forward pass. This is
useful for GPU parallelization. The batch DICE loss, if used, must be computed over the whole batch. Therefore, in a
naive implementation, all softmax outputs must be copied to a single GPU which will then
do the loss computation all by itself. In the context of 3D Segmentation, this results in a lot of overhead AND
is inefficient because the DICE computation is also kinda expensive (Think 8 GPUs with a result of shape
2x4x128x128x128 each.). The DICE is a global metric, but its parts can be computed locally (TP, FP, FN). Thus,
this implementation will compute all the parts of the loss function in the forward pass (and thus in a
parallelized way). The results are very small (batch_size x num_classes for TP, FN and FP, respectively; scalar for CE) and
copied easily. Also the final steps of the loss function (computing batch dice and average CE values) are easy
and very quick on the one GPU they need to run on. BAM.
final_nonlin is lambda x:x here!
"""
super(Generic_UNet_DP, self).__init__(
input_channels, base_num_features, num_classes, num_pool,
num_conv_per_stage, feat_map_mul_on_downscale, conv_op, norm_op,
norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin,
nonlin_kwargs, deep_supervision, dropout_in_localization,
lambda x: x, weightInitializer, pool_op_kernel_sizes,
conv_kernel_sizes, upscale_logits, convolutional_pooling,
convolutional_upsampling, max_num_features)
self.ce_loss = CrossentropyND()
def on_epoch_end(self):
"""
overwrite patient-based early stopping. Always run to 1000 epochs
:return:
"""
super().on_epoch_end()
continue_training = self.epoch < self.max_num_epochs
# it can rarely happen that the momentum of nnUNetTrainerV2 is too high for some dataset. If at epoch 100 the
# estimated validation Dice is still 0 then we reduce the momentum from 0.99 to 0.95
if self.epoch == 100:
if self.all_val_eval_metrics[-1] == 0:
self.optimizer.param_groups[0]["momentum"] = 0.95
self.network.apply(InitWeights_He(1e-2))
self.print_to_log_file("At epoch 100, the mean foreground Dice was 0. This can be caused by a too "
"high momentum. High momentum (0.99) is good for datasets where it works, but "
"sometimes causes issues such as this one. Momentum has now been reduced to "
"0.95 and network weights have been reinitialized")
return continue_training
def initialize_network(self):
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.ReLU
net_nonlin_kwargs = {'inplace': True}
self.network = Generic_UNet(self.num_input_channels,
self.base_num_features,
self.num_classes,
len(self.net_num_pool_op_kernel_sizes),
self.conv_per_stage,
2,
conv_op,
norm_op,
norm_op_kwargs,
dropout_op,
dropout_op_kwargs,
net_nonlin,
net_nonlin_kwargs,
True,
False,
lambda x: x,
InitWeights_He(0),
self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes,
False,
True,
True,
basic_block=ConvDropoutNonlinNorm)
if torch.cuda.is_available():
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def initialize_network(self):
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = Identity
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = Identity
norm_op_kwargs = {}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
self.network = Generic_UNet(
self.num_input_channels, self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes), self.conv_per_stage, 2,
conv_op, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, True, False, lambda x: x,
InitWeights_He(1e-2), self.net_num_pool_op_kernel_sizes,
self.net_conv_kernel_sizes, False, True, True)
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def initialize_network(self):
assert self.threeD is False
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
# print("sawNetTrainerMultiOpts.py:62", self.net_num_pool_op_kernel_sizes)
self.network = SAWNet(self.num_input_channels, self.base_num_features, self.num_classes,
len(self.net_num_pool_op_kernel_sizes),
self.conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs, # ds below
net_nonlin, net_nonlin_kwargs, False, False, softmax_helper, InitWeights_He(1e-2),
self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True, True)
if torch.cuda.is_available():
self.network.cuda()
def __init__(self,
input_channels,
base_num_features,
num_classes,
num_pool,
num_conv_per_stage=2,
feat_map_mul_on_downscale=2,
conv_op=nn.Conv2d,
norm_op=nn.BatchNorm2d,
norm_op_kwargs=None,
dropout_op=nn.Dropout2d,
dropout_op_kwargs=None,
nonlin=nn.LeakyReLU,
nonlin_kwargs=None,
deep_supervision=True,
dropout_in_localization=False,
final_nonlin=softmax_helper,
weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=None,
conv_kernel_sizes=None,
upscale_logits=False,
convolutional_pooling=False,
convolutional_upsampling=False,
max_num_features=None,
basic_block=ConvDropoutNormNonlin,
seg_output_use_bias=False):
"""
basically more flexible than v1, architecture is the same
Does this look complicated? Nah bro. Functionality > usability
This does everything you need, including world peace.
Questions? -> [email protected]
"""
super(Generic_UNet_attention, self).__init__()
self.convolutional_upsampling = convolutional_upsampling
self.convolutional_pooling = convolutional_pooling
self.upscale_logits = upscale_logits
if nonlin_kwargs is None:
nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
if dropout_op_kwargs is None:
dropout_op_kwargs = {'p': 0.5, 'inplace': True}
if norm_op_kwargs is None:
norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}
self.nonlin = nonlin
self.nonlin_kwargs = nonlin_kwargs
self.dropout_op_kwargs = dropout_op_kwargs
self.norm_op_kwargs = norm_op_kwargs
self.weightInitializer = weightInitializer
self.conv_op = conv_op
self.norm_op = norm_op
self.dropout_op = dropout_op
self.num_classes = num_classes
self.final_nonlin = final_nonlin
self._deep_supervision = deep_supervision
self.do_ds = deep_supervision
if conv_op == nn.Conv2d:
upsample_mode = 'bilinear'
pool_op = nn.MaxPool2d
transpconv = nn.ConvTranspose2d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
elif conv_op == nn.Conv3d:
upsample_mode = 'trilinear'
pool_op = nn.MaxPool3d
transpconv = nn.ConvTranspose3d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
else:
raise ValueError(
"unknown convolution dimensionality, conv op: %s" %
str(conv_op))
self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes,
0,
dtype=np.int64)
self.pool_op_kernel_sizes = pool_op_kernel_sizes
self.conv_kernel_sizes = conv_kernel_sizes
self.conv_pad_sizes = []
for krnl in self.conv_kernel_sizes:
self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])
if max_num_features is None:
if self.conv_op == nn.Conv3d:
self.max_num_features = self.MAX_NUM_FILTERS_3D
else:
self.max_num_features = self.MAX_FILTERS_2D
else:
self.max_num_features = max_num_features
self.conv_blocks_context = []
self.conv_blocks_localization = []
self.td = []
self.tu = []
self.att = []
self.seg_outputs = []
output_features = base_num_features
input_features = input_channels
for d in range(num_pool):
# determine the first stride
if d != 0 and self.convolutional_pooling:
first_stride = pool_op_kernel_sizes[d - 1]
else:
first_stride = None
self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
# add convolutions
self.conv_blocks_context.append(
StackedConvLayers(input_features,
output_features,
num_conv_per_stage,
self.conv_op,
self.conv_kwargs,
self.norm_op,
self.norm_op_kwargs,
self.dropout_op,
self.dropout_op_kwargs,
self.nonlin,
self.nonlin_kwargs,
first_stride,
basic_block=basic_block))
if not self.convolutional_pooling:
self.td.append(pool_op(pool_op_kernel_sizes[d]))
input_features = output_features
output_features = int(
np.round(output_features * feat_map_mul_on_downscale))
output_features = min(output_features, self.max_num_features)
# now the bottleneck.
# determine the first stride
if self.convolutional_pooling:
first_stride = pool_op_kernel_sizes[-1]
else:
first_stride = None
# the output of the last conv must match the number of features from the skip connection if we are not using
# convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
# done by the transposed conv
if self.convolutional_upsampling:
final_num_features = output_features
else:
final_num_features = self.conv_blocks_context[-1].output_channels
self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]
self.conv_blocks_context.append(
nn.Sequential(
StackedConvLayers(input_features,
output_features,
num_conv_per_stage - 1,
self.conv_op,
self.conv_kwargs,
self.norm_op,
self.norm_op_kwargs,
self.dropout_op,
self.dropout_op_kwargs,
self.nonlin,
self.nonlin_kwargs,
first_stride,
basic_block=basic_block),
StackedConvLayers(output_features,
final_num_features,
1,
self.conv_op,
self.conv_kwargs,
self.norm_op,
self.norm_op_kwargs,
self.dropout_op,
self.dropout_op_kwargs,
self.nonlin,
self.nonlin_kwargs,
basic_block=basic_block)))
# if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
if not dropout_in_localization:
old_dropout_p = self.dropout_op_kwargs['p']
self.dropout_op_kwargs['p'] = 0.0
# now lets build the localization pathway
for u in range(num_pool):
nfeatures_from_down = final_num_features
nfeatures_from_skip = self.conv_blocks_context[-(
2 + u
)].output_channels # self.conv_blocks_context[-1] is bottleneck, so start with -2
n_features_after_tu_and_concat = nfeatures_from_skip * 2
# the first conv reduces the number of features to match those of skip
# the following convs work on that number of features
# if not convolutional upsampling then the final conv reduces the num of features again
if u != num_pool - 1 and not self.convolutional_upsampling:
final_num_features = self.conv_blocks_context[-(
3 + u)].output_channels
else:
final_num_features = nfeatures_from_skip
if not self.convolutional_upsampling:
self.tu.append(
Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)],
mode=upsample_mode))
else:
self.tu.append(
transpconv(nfeatures_from_down,
nfeatures_from_skip,
pool_op_kernel_sizes[-(u + 1)],
pool_op_kernel_sizes[-(u + 1)],
bias=False))
self.att.append(
Attention(nfeatures_from_skip, nfeatures_from_skip // 2,
self.conv_op, self.conv_kwargs, self.norm_op,
self.norm_op_kwargs, self.nonlin,
self.nonlin_kwargs))
self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]
self.conv_blocks_localization.append(
nn.Sequential(
StackedConvLayers(n_features_after_tu_and_concat,
nfeatures_from_skip,
num_conv_per_stage - 1,
self.conv_op,
self.conv_kwargs,
self.norm_op,
self.norm_op_kwargs,
self.dropout_op,
self.dropout_op_kwargs,
self.nonlin,
self.nonlin_kwargs,
basic_block=basic_block),
StackedConvLayers(nfeatures_from_skip,
final_num_features,
1,
self.conv_op,
self.conv_kwargs,
self.norm_op,
self.norm_op_kwargs,
self.dropout_op,
self.dropout_op_kwargs,
self.nonlin,
self.nonlin_kwargs,
basic_block=basic_block)))
for ds in range(len(self.conv_blocks_localization)):
self.seg_outputs.append(
conv_op(self.conv_blocks_localization[ds][-1].output_channels,
num_classes, 1, 1, 0, 1, 1, seg_output_use_bias))
self.upscale_logits_ops = []
cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes),
axis=0)[::-1]
for usl in range(num_pool - 1):
if self.upscale_logits:
self.upscale_logits_ops.append(
Upsample(scale_factor=tuple(
[int(i) for i in cum_upsample[usl + 1]]),
mode=upsample_mode))
else:
self.upscale_logits_ops.append(lambda x: x)
if not dropout_in_localization:
self.dropout_op_kwargs['p'] = old_dropout_p
# register all modules properly
self.conv_blocks_localization = nn.ModuleList(
self.conv_blocks_localization)
self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
self.td = nn.ModuleList(self.td)
self.tu = nn.ModuleList(self.tu)
self.att = nn.ModuleList(self.att)
self.seg_outputs = nn.ModuleList(self.seg_outputs)
if self.upscale_logits:
self.upscale_logits_ops = nn.ModuleList(
self.upscale_logits_ops
) # lambda x:x is not a Module so we need to distinguish here
if self.weightInitializer is not None:
self.apply(self.weightInitializer)
def __init__(self,
input_channels,
base_num_features,
num_classes,
num_pool,
num_conv_per_stage=2,
feat_map_mul_on_downscale=2,
conv_op=nn.Conv2d,
norm_op=nn.BatchNorm2d,
norm_op_kwargs=None,
dropout_op=nn.Dropout2d,
dropout_op_kwargs=None,
nonlin=nn.LeakyReLU,
nonlin_kwargs=None,
deep_supervision=True,
dropout_in_localization=False,
final_nonlin=softmax_helper,
weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=None,
conv_kernel_sizes=None,
upscale_logits=False,
convolutional_pooling=False,
convolutional_upsampling=False,
max_num_features=None,
basic_block=ConvDropoutNormNonlin,
seg_output_use_bias=False):
super(SAWNet, self).__init__(
input_channels, base_num_features, num_classes, num_pool,
num_conv_per_stage, feat_map_mul_on_downscale, conv_op, norm_op,
norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin,
nonlin_kwargs, deep_supervision, dropout_in_localization,
final_nonlin, weightInitializer, pool_op_kernel_sizes,
conv_kernel_sizes, upscale_logits, convolutional_pooling,
convolutional_upsampling, max_num_features, basic_block,
seg_output_use_bias)
self.conv_blocks_w = []
self.tuw = []
self.w_outputs = []
upsample_mode = 'bilinear'
transpconv = nn.ConvTranspose2d
output_features = base_num_features
input_features = input_channels
print("num_pool:", num_pool)
print("input_features:", input_features)
for d in range(num_pool):
input_features = output_features
output_features = int(
np.round(output_features * feat_map_mul_on_downscale))
output_features = min(output_features, self.max_num_features)
print("{}: {}".format(d, input_features))
self.sau = SAUnit(output_features)
if self.convolutional_upsampling:
final_num_features = output_features # base_num_features
else:
final_num_features = self.conv_blocks_context[-1].output_channels
for u in range(num_pool):
nfeatures_from_down = final_num_features
nfeatures_from_skip = self.conv_blocks_context[-(
2 + u)].output_channels
n_features_after_tu_and_concat = nfeatures_from_skip * 2
# the first conv reduces the number of features to match those of skip
# the following convs work on that number of features
# if not convolutional upsampling then the final conv reduces the num of features again
if u != num_pool - 1 and not self.convolutional_upsampling:
final_num_features = self.conv_blocks_context[-(
3 + u)].output_channels
else:
final_num_features = nfeatures_from_skip
if not self.convolutional_upsampling:
self.tuw.append(
Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)],
mode=upsample_mode))
else:
# print("u: {} nfeatures_from_down: {} nfeatures_from_skip: {}".format(u, nfeatures_from_down,
# nfeatures_from_skip))
self.tuw.append(
transpconv(nfeatures_from_down,
nfeatures_from_skip,
pool_op_kernel_sizes[-(u + 1)],
pool_op_kernel_sizes[-(u + 1)],
bias=False))
self.conv_blocks_w.append(
nn.Sequential(
StackedConvLayers(n_features_after_tu_and_concat,
nfeatures_from_skip,
num_conv_per_stage - 1,
self.conv_op,
self.conv_kwargs,
self.norm_op,
self.norm_op_kwargs,
self.dropout_op,
self.dropout_op_kwargs,
self.nonlin,
self.nonlin_kwargs,
basic_block=basic_block),
StackedConvLayers(nfeatures_from_skip,
final_num_features,
1,
self.conv_op,
self.conv_kwargs,
self.norm_op,
self.norm_op_kwargs,
self.dropout_op,
self.dropout_op_kwargs,
self.nonlin,
self.nonlin_kwargs,
basic_block=basic_block)))
for ds in range(len(self.conv_blocks_w)):
# todo: relu?
self.w_outputs.append(
conv_op(self.conv_blocks_w[ds][-1].output_channels, 1, 1, 1, 0,
1, 1, seg_output_use_bias))
self.conv_blocks_w = nn.ModuleList(self.conv_blocks_w)
self.tuw = nn.ModuleList(self.tuw)
self.w_outputs = nn.ModuleList(self.w_outputs)
# self.final_conv = nn.Conv2d(32, 1, (1, 1))
if self.weightInitializer is not None:
self.apply(self.weightInitializer)
def __init__(self, steps,multiplier,input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage=2,
feat_map_mul_on_downscale=2, conv_op=nn.Conv2d,
norm_op=nn.BatchNorm2d, norm_op_kwargs=None,
dropout_op=nn.Dropout2d, dropout_op_kwargs=None,
nonlin=nn.LeakyReLU, nonlin_kwargs=None, deep_supervision=True, dropout_in_localization=False,
final_nonlin=softmax_helper, weightInitializer=InitWeights_He(1e-2), pool_op_kernel_sizes=None,
conv_kernel_sizes=None,
upscale_logits=False, convolutional_pooling=False, convolutional_upsampling=False,
max_num_features=None, basic_block=ConvDropoutNormNonlin,
seg_output_use_bias=False):
super(Nas_UNet, self).__init__()
self.convolutional_upsampling = convolutional_upsampling
self.convolutional_pooling = convolutional_pooling
self.upscale_logits = upscale_logits
self._steps=steps
self._multiplier=multiplier
if nonlin_kwargs is None:
nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
if dropout_op_kwargs is None:
dropout_op_kwargs = {'p': 0.5, 'inplace': True}
if norm_op_kwargs is None:
norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}
self.nonlin = nonlin
self.nonlin_kwargs = nonlin_kwargs
self.dropout_op_kwargs = dropout_op_kwargs
self.norm_op_kwargs = norm_op_kwargs
self.weightInitializer = weightInitializer
self.conv_op = conv_op
self.norm_op = norm_op
self.dropout_op = dropout_op
self.num_classes = num_classes
self.final_nonlin = final_nonlin
self._deep_supervision = deep_supervision
self.do_ds = deep_supervision
if conv_op == nn.Conv2d:
upsample_mode = 'bilinear'
pool_op = nn.MaxPool2d
transpconv = nn.ConvTranspose2d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
elif conv_op == nn.Conv3d:
upsample_mode = 'trilinear'
pool_op = nn.MaxPool3d
transpconv = nn.ConvTranspose3d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
else:
raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))
self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
self.pool_op_kernel_sizes = pool_op_kernel_sizes
self.conv_kernel_sizes = conv_kernel_sizes
self.conv_pad_sizes = []
for krnl in self.conv_kernel_sizes:
self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])
if max_num_features is None:
if self.conv_op == nn.Conv3d:
self.max_num_features = self.MAX_NUM_FILTERS_3D
else:
self.max_num_features = self.MAX_FILTERS_2D
else:
self.max_num_features = max_num_features
self.conv_blocks_context = []
self.conv_blocks_localization = []
self.td = []
self.tu = []
self.seg_outputs = []
output_features = base_num_features
input_features = input_channels
################ PART1
for d in range(2): # 保留前两个conv_blocks_context
# determine the first stride
if d != 0 and self.convolutional_pooling:
first_stride = pool_op_kernel_sizes[d - 1]
else:
first_stride = None
self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
# add convolutions
#print('input_features',input_features)
#print('output_features1',output_features)
self.conv_blocks_context.append(StackedConvLayers(input_features, output_features, num_conv_per_stage,
self.conv_op, self.conv_kwargs, self.norm_op,
self.norm_op_kwargs, self.dropout_op,
self.dropout_op_kwargs, self.nonlin, self.nonlin_kwargs,
first_stride, basic_block=basic_block))
if not self.convolutional_pooling:
self.td.append(pool_op(pool_op_kernel_sizes[d]))
input_features = output_features
#print('output_features2',output_features)
output_features = int(np.round(output_features * feat_map_mul_on_downscale))
# print('output_features3',output_features)
output_features = min(output_features, self.max_num_features)
# print('output_features4',output_features)
# print('output_features',output_features)
#exit(0)
################ PART2
# darts-nas，reduction cell
# output_features
self.reduction_cells = []
reduction=True
C_in, C_out=input_features,input_features
for i in range(3):
C_out=C_out*2
C_out=min(C_out,self.max_num_features)
reduction_cell = Cell(self._steps,self._multiplier,C_in//2,C_in,C_out,reduction,i)
self.reduction_cells.append(reduction_cell)
C_in = C_out
input_features = C_out
################ PART3
output_features = int(np.round(C_out * feat_map_mul_on_downscale))
output_features = min(output_features, self.max_num_features)
# now the bottleneck.
# determine the first stride
if self.convolutional_pooling:
first_stride = pool_op_kernel_sizes[-1]
else:
first_stride = None
# the output of the last conv must match the number of features from the skip connection if we are not using
# convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
# done by the transposed conv
if self.convolutional_upsampling:
final_num_features = output_features
else:
final_num_features = self.conv_blocks_context[-1].output_channels
self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]
self.conv_blocks_context.append(nn.Sequential(
StackedConvLayers(input_features, output_features, num_conv_per_stage - 1, self.conv_op, self.conv_kwargs,
self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, self.nonlin,
self.nonlin_kwargs, first_stride, basic_block=basic_block),
StackedConvLayers(output_features, final_num_features, 1, self.conv_op, self.conv_kwargs,
self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, self.nonlin,
self.nonlin_kwargs, basic_block=basic_block)))
# if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
if not dropout_in_localization:
old_dropout_p = self.dropout_op_kwargs['p']
self.dropout_op_kwargs['p'] = 0.0
################ PART4
##### normal_cell
# darts-nas，normal cell
# output_features
self.normal_cells = []
reduction=False
# C_in, C_out=final_num_features,final_num_features
C_in0=self.reduction_cells[2].output_channels
for i in range(3): # 0 1 2
C_out=self.reduction_cells[2-i].output_channels
C_in=C_out*2
normal_cell = Cell(self._steps,self._multiplier, C_in0,C_in,C_out,reduction,i)
self.normal_cells.append(normal_cell)
C_in0=C_in
# for tu
if not self.convolutional_upsampling:
self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
else:
self.tu.append(transpconv(final_num_features, C_out, pool_op_kernel_sizes[-(i + 1)],
pool_op_kernel_sizes[-(i + 1)], bias=False))
final_num_features=C_out
self._initialize_alphas()
################ PART5
# now lets build the localization pathway
for u in range(3,5): # 3 4
# nfeatures_from_down = final_num_features
nfeatures_from_skip = self.conv_blocks_context[4-u].output_channels
n_features_after_tu_and_concat = nfeatures_from_skip * 2
# u!= 4
# if u != num_pool - 1 and not self.convolutional_upsampling:
# final_num_features = self.conv_blocks_context[4-u].output_channels
# else:
# final_num_features = nfeatures_from_skip
if not self.convolutional_upsampling:
self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
else:
self.tu.append(transpconv(n_features_after_tu_and_concat, nfeatures_from_skip, pool_op_kernel_sizes[-(u + 1)],
pool_op_kernel_sizes[-(u + 1)], bias=False))
self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[- (u + 1)]
self.conv_kwargs['padding'] = self.conv_pad_sizes[- (u + 1)]
self.conv_blocks_localization.append(nn.Sequential(
StackedConvLayers(n_features_after_tu_and_concat, nfeatures_from_skip, num_conv_per_stage - 1,
self.conv_op, self.conv_kwargs, self.norm_op, self.norm_op_kwargs, self.dropout_op,
self.dropout_op_kwargs, self.nonlin, self.nonlin_kwargs, basic_block=basic_block),
StackedConvLayers(nfeatures_from_skip, nfeatures_from_skip, 1, self.conv_op, self.conv_kwargs,
self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs,
self.nonlin, self.nonlin_kwargs, basic_block=basic_block)
))
for ds in range(len(self.normal_cells)):
self.seg_outputs.append(conv_op(self.normal_cells[ds].output_channels, num_classes,
1, 1, 0, 1, 1, seg_output_use_bias))
for ds in range(len(self.conv_blocks_localization)):
self.seg_outputs.append(conv_op(self.conv_blocks_localization[ds][-1].output_channels, num_classes,
1, 1, 0, 1, 1, seg_output_use_bias))
self.upscale_logits_ops = []
cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
for usl in range(num_pool - 1):
if self.upscale_logits:
self.upscale_logits_ops.append(Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]),
mode=upsample_mode))
else:
self.upscale_logits_ops.append(lambda x: x)
if not dropout_in_localization:
self.dropout_op_kwargs['p'] = old_dropout_p
# register all modules properly
self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
self.td = nn.ModuleList(self.td)
self.tu = nn.ModuleList(self.tu)
self.normal_cells=nn.ModuleList(self.normal_cells)
self.reduction_cells=nn.ModuleList(self.reduction_cells)
self.seg_outputs = nn.ModuleList(self.seg_outputs)
if self.upscale_logits:
self.upscale_logits_ops = nn.ModuleList(
self.upscale_logits_ops) # lambda x:x is not a Module so we need to distinguish here
if self.weightInitializer is not None:
self.apply(self.weightInitializer)
def __init__(self,
input_channels,
base_num_features,
num_classes,
num_pool,
num_conv_per_stage=2,
feat_map_mul_on_downscale=2,
conv_op=nn.Conv2d,
norm_op=nn.BatchNorm2d,
norm_op_kwargs=None,
dropout_op=nn.Dropout2d,
dropout_op_kwargs=None,
nonlin=nn.LeakyReLU,
nonlin_kwargs=None,
deep_supervision=True,
dropout_in_localization=False,
final_nonlin=softmax_helper,
weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=None,
conv_kernel_sizes=None,
upscale_logits=False,
convolutional_pooling=False,
convolutional_upsampling=False,
max_num_features=None):
"""
basically more flexible than v1, architecture is the same
Does this look complicated? Nah bro. Functionality > usability
This does everything you need, including world peace.
Questions? -> [email protected]
"""
# conv_op = OctConv_3D
super(Generic_UNet, self).__init__()
self.convolutional_upsampling = convolutional_upsampling
self.convolutional_pooling = convolutional_pooling
self.upscale_logits = upscale_logits
if nonlin_kwargs is None:
nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
if dropout_op_kwargs is None:
dropout_op_kwargs = {'p': 0.5, 'inplace': True}
if norm_op_kwargs is None:
norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}
self.nonlin = nonlin
self.nonlin_kwargs = nonlin_kwargs
self.dropout_op_kwargs = dropout_op_kwargs
self.norm_op_kwargs = norm_op_kwargs
self.weightInitializer = weightInitializer
self.conv_op = conv_op
self.norm_op = norm_op
self.dropout_op = dropout_op
self.num_classes = num_classes
self.final_nonlin = final_nonlin
self.do_ds = deep_supervision
if conv_op == nn.Conv2d:
upsample_mode = 'bilinear'
pool_op = nn.MaxPool2d
transpconv = nn.ConvTranspose2d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
elif conv_op == nn.Conv3d:
upsample_mode = 'trilinear'
pool_op = nn.MaxPool3d
transpconv = nn.ConvTranspose3d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
elif conv_op == OctConv_3D:
upsample_mode = 'trilinear'
pool_op = nn.MaxPool3d
transpconv = nn.ConvTranspose3d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
else:
raise ValueError(
"unknown convolution dimensionality, conv op: %s" %
str(conv_op))
self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes,
0,
dtype=np.int64)
self.pool_op_kernel_sizes = pool_op_kernel_sizes
self.conv_kernel_sizes = conv_kernel_sizes
self.conv_pad_sizes = []
for krnl in self.conv_kernel_sizes:
self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])
if max_num_features is None:
if self.conv_op == nn.Conv3d:
self.max_num_features = self.MAX_NUM_FILTERS_3D
else:
self.max_num_features = self.MAX_FILTERS_2D
else:
self.max_num_features = max_num_features
self.conv_blocks_context = []
self.conv_blocks_localization = []
self.td = []
self.tu = []
self.seg_outputs = []
output_features = base_num_features
input_features = input_channels
# model = UNet_Nested(in_channels=input_channels, n_classes=num_classes)
# model = model.cuda()
# if torch.cuda.device_count() > 1:
# # torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=0,
# # world_size=1)
# # process_group = torch.distributed.new_group()
# # model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
# model = DistributedDataParallel(model)
# print("Let's use", torch.cuda.device_count(), "GPUs!")
# # 5) 封装
# # model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True)
self.UNet_Nested = UNet_Nested(in_channels=input_channels,
n_classes=num_classes)
def __init__(self,
img_size,
input_channels,
base_num_features,
num_classes,
num_pool,
num_conv_per_stage=2,
feat_map_mul_on_downscale=2,
conv_op=nn.Conv2d,
norm_op=nn.BatchNorm2d,
norm_op_kwargs=None,
dropout_op=nn.Dropout2d,
dropout_op_kwargs=None,
nonlin=nn.LeakyReLU,
nonlin_kwargs=None,
deep_supervision=False,
dropout_in_localization=False,
final_nonlin=softmax_helper,
weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=None,
conv_kernel_sizes=None,
upscale_logits=False,
convolutional_pooling=False,
convolutional_upsampling=False,
max_num_features=None,
seg_output_use_bias=False,
zero_head=False,
vis=False):
super(VisionTransformer, self).__init__()
DEFAULT_BATCH_SIZE_3D = 4
DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
SPACING_FACTOR_BETWEEN_STAGES = 2
BASE_NUM_FEATURES_3D = 30
MAX_NUMPOOL_3D = 999
MAX_NUM_FILTERS_3D = 320
DEFAULT_PATCH_SIZE_2D = (256, 256)
BASE_NUM_FEATURES_2D = 30
DEFAULT_BATCH_SIZE_2D = 50
MAX_NUMPOOL_2D = 999
MAX_FILTERS_2D = 480
use_this_for_batch_size_computation_2D = 19739648
use_this_for_batch_size_computation_3D = 520000000 # 505789440
config = ml_collections.ConfigDict()
config.patches = ml_collections.ConfigDict({'size': (16, 16, 16)})
config.hidden_size = 768
config.transformer = ml_collections.ConfigDict()
config.transformer.mlp_dim = 3072
config.transformer.num_heads = 12
config.transformer.num_layers = 12
config.transformer.attention_dropout_rate = 0.0
config.transformer.dropout_rate = 0.1
config.representation_size = None
config.resnet_pretrained_path = None
config.patch_size = 16
config.n_classes = 1
config.patches.grid = (4, 16, 16)
config.resnet = ml_collections.ConfigDict()
config.resnet.num_layers = (3, 4, 9)
config.resnet.width_factor = 1
config.classifier = 'seg'
config.pretrained_path = '../model/vit_checkpoint/imagenet21k/R50+ViT-B_16.npz'
config.decoder_channels = (256, 128, 64, 16)
config.skip_channels = [512, 256, 64, 16]
config.n_skip = 3
config.activation = 'softmax'
self.num_classes = num_classes
self.zero_head = zero_head
self.classifier = config.classifier
self.transformer = Transformer(config, img_size, vis)
self.decoder = DecoderCup(config)
self.segmentation_head = SegmentationHead(
in_channels=config['decoder_channels'][-1],
out_channels=config['n_classes'],
kernel_size=3,
)
self.config = config
self.convolutional_upsampling = convolutional_upsampling # True
self.convolutional_pooling = convolutional_pooling # True
self.upscale_logits = upscale_logits # False
if nonlin_kwargs is None:
nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True} #
if dropout_op_kwargs is None:
dropout_op_kwargs = {'p': 0.5, 'inplace': True}
if norm_op_kwargs is None:
norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}
self.nonlin = nonlin # nn.LeakyReLU
self.nonlin_kwargs = nonlin_kwargs # {'negative_slope': 1e-2, 'inplace': True}
self.dropout_op_kwargs = dropout_op_kwargs # {'p': 0, 'inplace': True}
self.norm_op_kwargs = norm_op_kwargs # {'eps': 1e-5, 'affine': True}
self.weightInitializer = weightInitializer # InitWeights_He(1e-2)
self.conv_op = conv_op # nn.Conv3d
self.norm_op = norm_op # nn.InstanceNorm3d
self.dropout_op = dropout_op # nn.Dropout3d
self.num_classes = num_classes
self.final_nonlin = final_nonlin
self._deep_supervision = deep_supervision
self.do_ds = deep_supervision
if conv_op == nn.Conv2d:
upsample_mode = 'bilinear'
pool_op = nn.MaxPool2d
transpconv = nn.ConvTranspose2d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
elif conv_op == nn.Conv3d:
upsample_mode = 'trilinear'
pool_op = nn.MaxPool3d
transpconv = nn.ConvTranspose3d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
else:
raise ValueError(
"unknown convolution dimensionality, conv op: %s" %
str(conv_op))
self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes,
0,
dtype=np.int64)
self.pool_op_kernel_sizes = pool_op_kernel_sizes
self.conv_kernel_sizes = conv_kernel_sizes
self.conv_pad_sizes = []
for krnl in self.conv_kernel_sizes:
self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])
if max_num_features is None:
if self.conv_op == nn.Conv3d:
self.max_num_features = self.MAX_NUM_FILTERS_3D
else:
self.max_num_features = self.MAX_FILTERS_2D
else:
self.max_num_features = max_num_features
if self.weightInitializer is not None:
self.apply(self.weightInitializer)
def initialize_network_optimizer_and_scheduler(self):
"""
This is specific to the U-Net and must be adapted for other network architectures
:return:
"""
#self.print_to_log_file(self.net_num_pool_op_kernel_sizes)
#self.print_to_log_file(self.net_conv_kernel_sizes)
print('Has entered into the initialize_network_optimizer_and_scheduler！')
net_numpool = len(self.net_num_pool_op_kernel_sizes)
if self.threeD:
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
else:
conv_op = nn.Conv2d
dropout_op = nn.Dropout2d
norm_op = nn.InstanceNorm2d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
torch.cuda.set_device(0)
self.do_supervision=False
self.network = Generic_UNet(self.num_input_channels, self.base_num_features, self.num_classes, net_numpool,
2, 2, conv_op, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, self.do_supervision, False, lambda x: x, InitWeights_He(1e-2),
self.net_num_pool_op_kernel_sizes, self.net_conv_kernel_sizes, False, True, True)
# print('self.network is:',self.network)
# Here!!! Modify it to be parallel model
# self.network=nn.DataParallel(self.network,device_ids=[0,1,2,3])
g=tw.draw_model(self.network,[8,1,128,128,128])
g.save('/cache/WenshuaiZhao/ProjectFiles/NNUnet/nnunet/model_structure.png')
print('draw model has been done!')
self.optimizer = torch.optim.Adam(self.network.parameters(), self.initial_lr, weight_decay=self.weight_decay, amsgrad=True)
# Here!!! optimizer is also need to be parallel
# self.optimizer=nn.DataParallel(self.optimizer,device_ids=[0,1,2,3])
self.lr_scheduler = lr_scheduler.ReduceLROnPlateau(self.optimizer, mode='min', factor=0.2, patience=self.lr_scheduler_patience,
verbose=True, threshold=self.lr_scheduler_eps, threshold_mode="abs")
self.network.cuda()
self.network.inference_apply_nonlin = softmax_helper
def __init__(self,
input_channels,
base_num_features,
num_classes,
num_pool,
num_conv_per_stage=2,
feat_map_mul_on_downscale=2,
conv_op=nn.Conv3d,
norm_op=nn.BatchNorm3d,
norm_op_kwargs=None,
dropout_op=nn.Dropout3d,
dropout_op_kwargs=None,
nonlin=nn.LeakyReLU,
nonlin_kwargs=None,
deep_supervision=False,
dropout_in_localization=False,
final_nonlin=softmax_helper,
weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=None,
conv_kernel_sizes=None,
upscale_logits=False,
convolutional_pooling=False,
convolutional_upsampling=False,
max_num_features=None,
basic_block=ConvDropoutNormNonlin,
seg_output_use_bias=False):
super(AttU_Net, self).__init__()
self.convolutional_upsampling = convolutional_upsampling # True
self.convolutional_pooling = convolutional_pooling # True
self.upscale_logits = upscale_logits # False
if nonlin_kwargs is None:
nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True} #
if dropout_op_kwargs is None:
dropout_op_kwargs = {'p': 0.5, 'inplace': True}
if norm_op_kwargs is None:
norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}
self.nonlin = nonlin # nn.LeakyReLU
self.nonlin_kwargs = nonlin_kwargs # {'negative_slope': 1e-2, 'inplace': True}
self.dropout_op_kwargs = dropout_op_kwargs # {'p': 0, 'inplace': True}
self.norm_op_kwargs = norm_op_kwargs # {'eps': 1e-5, 'affine': True}
self.weightInitializer = weightInitializer # InitWeights_He(1e-2)
self.conv_op = conv_op # nn.Conv3d
self.norm_op = norm_op # nn.InstanceNorm3d
self.dropout_op = dropout_op # nn.Dropout3d
self.num_classes = num_classes
self.final_nonlin = final_nonlin
self._deep_supervision = deep_supervision
self.do_ds = deep_supervision
if conv_op == nn.Conv2d:
upsample_mode = 'bilinear'
pool_op = nn.MaxPool2d
transpconv = nn.ConvTranspose2d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
elif conv_op == nn.Conv3d:
upsample_mode = 'trilinear'
pool_op = nn.MaxPool3d
transpconv = nn.ConvTranspose3d
if pool_op_kernel_sizes is None:
pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
if conv_kernel_sizes is None:
conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
else:
raise ValueError(
"unknown convolution dimensionality, conv op: %s" %
str(conv_op))
self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes,
0,
dtype=np.int64)
self.pool_op_kernel_sizes = pool_op_kernel_sizes
self.conv_kernel_sizes = conv_kernel_sizes
self.conv_pad_sizes = []
for krnl in self.conv_kernel_sizes:
self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])
if max_num_features is None:
if self.conv_op == nn.Conv3d:
self.max_num_features = self.MAX_NUM_FILTERS_3D
else:
self.max_num_features = self.MAX_FILTERS_2D
else:
self.max_num_features = max_num_features
if self.weightInitializer is not None:
self.apply(self.weightInitializer)
n1 = 16
filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]
self.Maxpool1 = nn.AvgPool3d(kernel_size=2, stride=2)
self.Maxpool2 = nn.AvgPool3d(kernel_size=2, stride=2)
self.Maxpool3 = nn.AvgPool3d(kernel_size=2, stride=2)
self.Maxpool4 = nn.AvgPool3d(kernel_size=2, stride=2)
self.Conv1 = conv_block(input_channels, filters[0])
self.Conv2 = conv_block(filters[0], filters[1])
self.Conv3 = conv_block(filters[1], filters[2])
self.Conv4 = conv_block(filters[2], filters[3])
self.Conv5 = conv_block(filters[3], filters[4])
self.Recons_up_conv1 = up_conv(filters[1], filters[0])
self.Recons_up_conv2 = up_conv(filters[2], filters[1])
self.transformer = Transformer()
self.Up5 = up_conv(filters[4], filters[3])
self.Att5 = Attention_block(F_g=filters[3],
F_l=filters[3],
F_int=filters[2])
self.Up_conv5 = conv_block(filters[4], filters[3])
self.Up4 = up_conv(filters[3], filters[2])
self.Att4 = Attention_block(F_g=filters[2],
F_l=filters[2],
F_int=filters[1])
self.Up_conv4 = conv_block(filters[3], filters[2])
self.Up3 = up_conv(filters[2], filters[1])
self.Att3 = Attention_block(F_g=filters[1],
F_l=filters[1],
F_int=filters[0])
self.Up_conv3 = conv_block(filters[2], filters[1])
self.Up2 = up_conv(filters[1], filters[0])
self.Att2 = Attention_block(F_g=filters[0], F_l=filters[0], F_int=32)
self.Up_conv2 = conv_block(filters[1], filters[0])
self.Conv = nn.Conv3d(filters[0],
num_classes,
kernel_size=1,
stride=1,
padding=0)
rates = (1, 6, 12, 18)
self.aspp1 = ASPP_module(16, 16, rate=rates[0])
self.aspp2 = ASPP_module(16, 16, rate=rates[1])
self.aspp3 = ASPP_module(16, 16, rate=rates[2])
self.aspp4 = ASPP_module(16, 16, rate=rates[3])
self.aspp_conv = nn.Conv3d(64, 64, 1)
self.aspp_gn = nn.GroupNorm(32, 64)
self.Out = nn.Conv3d(64,
num_classes,
kernel_size=1,
stride=1,
padding=0)
def time_train_iteration(model, inputs):
t1 = time.time()
preds = model(inputs)
t2 = time.time()
return t2-t1
if __name__ == '__main__':
from torch.optim import lr_scheduler
import time
model = Generic_UNet(input_channels=1, base_num_features=30, num_classes=2, num_pool=5, num_conv_per_stage=2,
feat_map_mul_on_downscale=2, conv_op=nn.Conv3d,
norm_op=nn.BatchNorm3d, norm_op_kwargs=None,
dropout_op=nn.Dropout3d, dropout_op_kwargs=None,
nonlin=nn.LeakyReLU, nonlin_kwargs=None, deep_supervision=True, dropout_in_localization=False,
final_nonlin=softmax_helper, weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=[[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]],
conv_kernel_sizes=[[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]],
upscale_logits=False, convolutional_pooling=True, convolutional_upsampling=True)
model.cuda()
inputs = torch.zeros([2, 1, 64, 128, 128]).cuda()
optimizer = torch.optim.Adam(model.parameters(), 3e-4, weight_decay=3e-5, amsgrad=True)
lr_scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, patience=30,
verbose=True, threshold=1e-3, threshold_mode="abs")
# try train itteration
print("***Try training***")
model.train()
rep = 100
t1 = time.time()
t = 0
net_numpool=len(net_num_pool_op_kernel_sizes)
conv_per_stage = 2
conv_op = nn.Conv3d
dropout_op = nn.Dropout3d
norm_op = nn.InstanceNorm3d
norm_op_kwargs = {'eps': 1e-5, 'affine': True}
dropout_op_kwargs = {'p': 0, 'inplace': True}
net_nonlin = nn.LeakyReLU
net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
##############################
model = Generic_UNet(num_input_channels, base_num_features, num_classes,
net_numpool,
conv_per_stage, 2, conv_op, norm_op, norm_op_kwargs, dropout_op,
dropout_op_kwargs,
net_nonlin, net_nonlin_kwargs, False, False, lambda x: x, InitWeights_He(1e-2),
net_num_pool_op_kernel_sizes, net_conv_kernel_sizes, False, True, True)
model.to(device)
summary(model, (1,64,64,32), batch_size=-1)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), config.lr, momentum=0.9, weight_decay=0.0, nesterov=False)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(config.patience * 0.6), gamma=0.5)
training_generator = generate_pair(x_train,config.batch_size,config)
validation_generator = generate_pair(x_valid, config.batch_size,config)
# to track the training loss as the model trains
train_losses = []
# to track the validation loss as the model trains
valid_losses = []
# to track the average training loss per epoch as the model trains
def test_calculate_metric(iter_nums):
if args.net == 'unet':
# timm-efficientnet performs slightly worse.
if not args.vis_mode:
backbone_type = re.sub("^eff", "efficientnet", args.backbone_type)
net = smp.Unet(backbone_type, classes=args.num_classes, encoder_weights='imagenet')
else:
net = VanillaUNet(n_channels=3, num_classes=args.num_classes)
elif args.net == 'unet-scratch':
# net = UNet(num_classes=args.num_classes)
net = VanillaUNet(n_channels=3, num_classes=args.num_classes,
use_polyformer=args.polyformer_mode,
num_modes=args.num_modes)
elif args.net == 'nestedunet':
net = NestedUNet(num_classes=args.num_classes)
elif args.net == 'unet3plus':
net = UNet_3Plus(num_classes=args.num_classes)
elif args.net == 'pranet':
net = PraNet(num_classes=args.num_classes - 1)
elif args.net == 'attunet':
net = AttU_Net(output_ch=args.num_classes)
elif args.net == 'attr2unet':
net = R2AttU_Net(output_ch=args.num_classes)
elif args.net == 'dunet':
net = DeformableUNet(n_channels=3, n_classes=args.num_classes)
elif args.net == 'setr':
# Install mmcv first:
# pip install mmcv-full==1.2.2 -f https://download.openmmlab.com/mmcv/dist/cu{Your CUDA Version}/torch{Your Pytorch Version}/index.html
# E.g.: pip install mmcv-full==1.2.2 -f https://download.openmmlab.com/mmcv/dist/cu102/torch1.7.1/index.html
from mmcv.utils import Config
sys.path.append("networks/setr")
from mmseg.models import build_segmentor
task2config = { 'refuge': 'SETR_PUP_288x288_10k_refuge_context_bs_4.py',
'polyp': 'SETR_PUP_320x320_10k_polyp_context_bs_4.py' }
setr_cfg = Config.fromfile("networks/setr/configs/SETR/{}".format(task2config[args.task_name]))
net = build_segmentor(setr_cfg.model, train_cfg=setr_cfg.train_cfg, test_cfg=setr_cfg.test_cfg)
# By default, net() calls forward_train(), which receives extra parameters, and returns losses.
# net.forward_dummy() receives/returns the traditional input/output.
# Relevant file: mmseg/models/segmentors/encoder_decoder.py
net.forward = net.forward_dummy
elif args.net == 'transunet':
transunet_config = TransUNet_CONFIGS[args.backbone_type]
transunet_config.n_classes = args.num_classes
if args.backbone_type.find('R50') != -1:
# The "patch" in TransUNet means grid-like patches of the input image.
# The "patch" in our code means the whole input image after cropping/resizing (part of the augmentation)
transunet_config.patches.grid = (int(args.patch_size[0] / transunet_config.patches.size[0]),
int(args.patch_size[1] / transunet_config.patches.size[1]))
net = TransUNet(transunet_config, img_size=args.patch_size, num_classes=args.num_classes)
elif args.net.startswith('deeplab'):
use_smp_deeplab = args.net.endswith('smp')
if use_smp_deeplab:
backbone_type = re.sub("^eff", "efficientnet", args.backbone_type)
net = smp.DeepLabV3Plus(backbone_type, classes=args.num_classes, encoder_weights='imagenet')
else:
model_name = args.net + "_" + args.backbone_type
model_map = {
'deeplabv3_resnet50': deeplab.deeplabv3_resnet50,
'deeplabv3plus_resnet50': deeplab.deeplabv3plus_resnet50,
'deeplabv3_resnet101': deeplab.deeplabv3_resnet101,
'deeplabv3plus_resnet101': deeplab.deeplabv3plus_resnet101,
'deeplabv3_mobilenet': deeplab.deeplabv3_mobilenet,
'deeplabv3plus_mobilenet': deeplab.deeplabv3plus_mobilenet
}
net = model_map[model_name](num_classes=args.num_classes, output_stride=8)
elif args.net == 'nnunet':
from nnunet.network_architecture.initialization import InitWeights_He
from nnunet.network_architecture.generic_UNet import Generic_UNet
net = Generic_UNet(
input_channels=3,
base_num_features=32,
num_classes=args.num_classes,
num_pool=7,
num_conv_per_stage=2,
feat_map_mul_on_downscale=2,
norm_op=nn.InstanceNorm2d,
norm_op_kwargs={'eps': 1e-05, 'affine': True},
dropout_op_kwargs={'p': 0, 'inplace': True},
nonlin_kwargs={'negative_slope': 0.01, 'inplace': True},
final_nonlin=(lambda x: x),
weightInitializer=InitWeights_He(1e-2),
pool_op_kernel_sizes=[[2, 2], [2, 2], [2, 2], [2, 2], [2, 2], [2, 2], [2, 2]],
conv_kernel_sizes=([[3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3]]),
upscale_logits=False,
convolutional_pooling=True,
convolutional_upsampling=True,
)
net.inference_apply_nonlin = (lambda x: F.softmax(x, 1))
elif args.net == 'segtran':
get_default(args, 'num_modes', default_settings, -1, [args.net, 'num_modes', args.in_fpn_layers])
set_segtran2d_config(args)
print(args)
net = Segtran2d(config)
else:
breakpoint()
net.cuda()
net.eval()
if args.robust_ref_cp_path:
refnet = copy.deepcopy(net)
print("Reference network created")
load_model(refnet, args, args.robust_ref_cp_path)
else:
refnet = None
# Currently colormap is used only for OCT task.
colormap = get_seg_colormap(args.num_classes, return_torch=True).cuda()
# prepred: pre-prediction. postpred: post-prediction.
task2mask_prepred = { 'refuge': refuge_map_mask, 'polyp': polyp_map_mask,
'oct': partial(index_to_onehot, num_classes=args.num_classes) }
task2mask_postpred = { 'refuge': refuge_inv_map_mask, 'polyp': polyp_inv_map_mask,
'oct': partial(onehot_inv_map, colormap=colormap) }
mask_prepred_mapping_func = task2mask_prepred[args.task_name]
mask_postpred_mapping_funcs = [ task2mask_postpred[args.task_name] ]
if args.do_remove_frag:
remove_frag = lambda segmap: remove_fragmentary_segs(segmap, 255)
mask_postpred_mapping_funcs.append(remove_frag)
if not args.checkpoint_dir:
if args.vis_mode is not None:
visualize_model(net, args.vis_mode, args.vis_layers, args.patch_size, db_test)
return
if args.eval_robustness:
eval_robustness(args, net, refnet, testloader, mask_prepred_mapping_func)
return
allcls_avg_metric = None
all_results = np.zeros((args.num_classes, len(iter_nums)))
for iter_idx, iter_num in enumerate(iter_nums):
if args.checkpoint_dir:
checkpoint_path = os.path.join(args.checkpoint_dir, 'iter_' + str(iter_num) + '.pth')
load_model(net, args, checkpoint_path)
if args.vis_mode is not None:
visualize_model(net, args.vis_mode, args.vis_layers, args.patch_size, db_test)
continue
if args.eval_robustness:
eval_robustness(args, net, refnet, testloader, mask_prepred_mapping_func)
continue
save_results = args.save_results and (not args.test_interp)
if save_results:
test_save_paths = []
test_save_dirs = []
test_save_dir_tmpl = "%s-%s-%s-%d" %(args.net, args.job_name, timestamp, iter_num)
for suffix in ("-soft", "-%.1f" %args.mask_thres):
test_save_dir = test_save_dir_tmpl + suffix
test_save_path = "../prediction/%s" %(test_save_dir)
if not os.path.exists(test_save_path):
os.makedirs(test_save_path)
test_save_dirs.append(test_save_dir)
test_save_paths.append(test_save_path)
else:
test_save_paths = None
test_save_dirs = None
if args.save_features_img_count > 0:
args.save_features_file_path = "%s-%s-feat-%s.pth" %(args.net, args.job_name, timestamp)
else:
args.save_features_file_path = None
allcls_avg_metric, allcls_metric_count = \
test_all_cases(net, testloader, task_name=args.task_name,
num_classes=args.num_classes,
mask_thres=args.mask_thres,
model_type=args.net,
orig_input_size=args.orig_input_size,
patch_size=args.patch_size,
stride=(args.orig_input_size[0] // 2, args.orig_input_size[1] // 2),
test_save_paths=test_save_paths,
out_origsize=args.out_origsize,
mask_prepred_mapping_func=mask_prepred_mapping_func,
mask_postpred_mapping_funcs=mask_postpred_mapping_funcs,
reload_mask=args.reload_mask,
test_interp=args.test_interp,
save_features_img_count=args.save_features_img_count,
save_features_file_path=args.save_features_file_path,
verbose=args.verbose_output)
print("Iter-%d scores on %d images:" %(iter_num, allcls_metric_count[0]))
dice_sum = 0
for cls in range(1, args.num_classes):
dice = allcls_avg_metric[cls-1]
print('class %d: dice = %.3f' %(cls, dice))
dice_sum += dice
all_results[cls, iter_idx] = dice
avg_dice = dice_sum / (args.num_classes - 1)
print("Average dice: %.3f" %avg_dice)
if args.net == 'segtran':
max_attn, avg_attn, clamp_count, call_count = \
[ segtran_shared.__dict__[v] for v in ('max_attn', 'avg_attn', 'clamp_count', 'call_count') ]
print("max_attn={:.2f}, avg_attn={:.2f}, clamp_count={}, call_count={}".format(
max_attn, avg_attn, clamp_count, call_count))
if save_results:
FNULL = open(os.devnull, 'w')
for pred_type, test_save_dir, test_save_path in zip(('soft', 'hard'), test_save_dirs, test_save_paths):
do_tar = subprocess.run(["tar", "cvf", "%s.tar" %test_save_dir, test_save_dir], cwd="../prediction",
stdout=FNULL, stderr=subprocess.STDOUT)
# print(do_tar)
print("{} tarball:\n{}.tar".format(pred_type, os.path.abspath(test_save_path)))
np.set_printoptions(precision=3, suppress=True)
print(all_results[1:])
return allcls_avg_metric
