model.load_state_dict(state['model_state_dict'])
optimizer_state_dict = state.get('optimizer_state_dict')
lr_sched_state_dict = state.get('lr_sched_state_dict')
if optimizer_state_dict is None:
logger.warning('optimizer_state_dict not found.')
if lr_sched_state_dict is None:
logger.warning('lr_sched_state_dict not found.')
elif isinstance(state, nn.Module):
logger.warning(_warning_str)
model = state
else:
raise ValueError(f'Can\'t load {pretrained}.')
# Transformations to be applied to samples before feeding them to the network
common_transforms = [
transforms.SqueezeTarget(dim=0), # Workaround for neuro_data_cdhw
transforms.Normalize(mean=dataset_mean, std=dataset_std)
]
train_transform = transforms.Compose(common_transforms + [
# transforms.RandomRotate2d(prob=0.9),
# transforms.RandomGrayAugment(channels=[0], prob=0.3),
# transforms.RandomGammaCorrection(gamma_std=0.25, gamma_min=0.25, prob=0.3),
# transforms.AdditiveGaussianNoise(sigma=0.1, channels=[0], prob=0.3),
])
valid_transform = transforms.Compose(common_transforms + [])
# Specify data set
aniso_factor = 2 # Anisotropy in z dimension. E.g. 2 means half resolution in z dimension.
common_data_kwargs = { # Common options for training and valid sets.
'aniso_factor': aniso_factor,
'patch_shape': (44, 88, 88),
model.load_state_dict(state['model_state_dict'])
optimizer_state_dict = state.get('optimizer_state_dict')
lr_sched_state_dict = state.get('lr_sched_state_dict')
if optimizer_state_dict is None:
logger.warning('optimizer_state_dict not found.')
if lr_sched_state_dict is None:
logger.warning('lr_sched_state_dict not found.')
elif isinstance(state, nn.Module):
logger.warning(_warning_str)
model = state
else:
raise ValueError(f'Can\'t load {pretrained}.')
# Transformations to be applied to samples before feeding them to the network
common_transforms = [
transforms.SqueezeTarget(dim=0),
]
train_transform = transforms.Compose(common_transforms + [
transforms.RandomFlip(ndim_spatial=3),
transforms.RandomGrayAugment(channels=[0], prob=0.3),
transforms.RandomGammaCorrection(gamma_std=0.25, gamma_min=0.25, prob=0.3),
transforms.AdditiveGaussianNoise(sigma=0.1, channels=[0], prob=0.3),
])
valid_transform = transforms.Compose(common_transforms + [])
# Specify data set
aniso_factor = 2 # Anisotropy in z dimension. E.g. 2 means half resolution in z dimension.
common_data_kwargs = { # Common options for training and valid sets.
'aniso_factor': aniso_factor,
'patch_shape': (48, 96, 96),
# 'offset': (8, 20, 20),
