def net(net_params, rtn_level=False):
"""Define the neural net"""
model_name = net_params['global']['model_name'].lower()
state_dict_path = ''
if model_name == 'unetsmall':
model = unet.UNetSmall(
net_params['global']['num_classes'],
net_params['global']['number_of_bands'],
net_params['models']['unetsmall']['dropout'],
net_params['models']['unetsmall']['probability'])
if net_params['models']['unetsmall']['pretrained']:
state_dict_path = net_params['models']['unetsmall']['pretrained']
elif model_name == 'unet':
model = unet.UNet(net_params['global']['num_classes'],
net_params['global']['number_of_bands'],
net_params['models']['unet']['dropout'],
net_params['models']['unet']['probability'])
if net_params['models']['unet']['pretrained']:
state_dict_path = net_params['models']['unet']['pretrained']
elif model_name == 'ternausnet':
model = TernausNet.ternausnet(
net_params['global']['num_classes'],
net_params['models']['ternausnet']['pretrained'])
elif model_name == 'checkpointed_unet':
model = checkpointed_unet.UNetSmall(
net_params['global']['num_classes'],
net_params['global']['number_of_bands'],
net_params['models']['unetsmall']['dropout'],
net_params['models']['unetsmall']['probability'])
if net_params['models']['unetsmall']['pretrained']:
state_dict_path = net_params['models']['unetsmall']['pretrained']
elif model_name == 'inception':
model = inception.Inception3(net_params['global']['num_classes'],
net_params['global']['number_of_bands'])
if net_params['models']['inception']['pretrained']:
state_dict_path = net_params['models']['inception']['pretrained']
else:
raise ValueError('The model name in the config.yaml is not defined.')
if rtn_level:
lvl = maxpool_level(model, net_params['global']['number_of_bands'],
256)
return model, state_dict_path, lvl['MaxPoolCount']
else:
return model, state_dict_path
def net(net_params, num_channels, inference=False):
"""Define the neural net"""
model_name = net_params['global']['model_name'].lower()
num_bands = int(net_params['global']['number_of_bands'])
msg = f'Number of bands specified incompatible with this model. Requires 3 band data.'
train_state_dict_path = get_key_def('state_dict_path',
net_params['training'], None)
pretrained = get_key_def('pretrained', net_params['training'],
True) if not inference else False
dropout = get_key_def('dropout', net_params['training'], False)
dropout_prob = get_key_def('dropout_prob', net_params['training'], 0.5)
if model_name == 'unetsmall':
model = unet.UNetSmall(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'unet':
model = unet.UNet(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'ternausnet':
assert num_bands == 3, msg
model = TernausNet.ternausnet(num_channels)
elif model_name == 'checkpointed_unet':
model = checkpointed_unet.UNetSmall(num_channels, num_bands, dropout,
dropout_prob)
elif model_name == 'inception':
model = inception.Inception3(num_channels, num_bands)
elif model_name == 'fcn_resnet101':
assert num_bands == 3, msg
model = models.segmentation.fcn_resnet101(pretrained=False,
progress=True,
num_classes=num_channels,
aux_loss=None)
elif model_name == 'deeplabv3_resnet101':
assert (num_bands == 3 or num_bands == 4), msg
if num_bands == 3:
print('Finetuning pretrained deeplabv3 with 3 bands')
model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True,
aux_loss=None)
model.classifier = common.DeepLabHead(2048, num_channels)
elif num_bands == 4:
print('Finetuning pretrained deeplabv3 with 4 bands')
model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True,
aux_loss=None)
conv1 = model.backbone._modules['conv1'].weight.detach().numpy()
depth = np.random.uniform(low=-1, high=1, size=(64, 1, 7, 7))
conv1 = np.append(conv1, depth, axis=1)
conv1 = torch.from_numpy(conv1).float()
model.backbone._modules['conv1'].weight = nn.Parameter(
conv1, requires_grad=True)
model.classifier = common.DeepLabHead(2048, num_channels)
else:
raise ValueError(
f'The model name {model_name} in the config.yaml is not defined.')
coordconv_convert = get_key_def('coordconv_convert', net_params['global'],
False)
if coordconv_convert:
centered = get_key_def('coordconv_centered', net_params['global'],
True)
normalized = get_key_def('coordconv_normalized', net_params['global'],
True)
noise = get_key_def('coordconv_noise', net_params['global'], None)
radius_channel = get_key_def('coordconv_radius_channel',
net_params['global'], False)
scale = get_key_def('coordconv_scale', net_params['global'], 1.0)
# note: this operation will not attempt to preserve already-loaded model parameters!
model = coordconv.swap_coordconv_layers(model,
centered=centered,
normalized=normalized,
noise=noise,
radius_channel=radius_channel,
scale=scale)
if inference:
state_dict_path = net_params['inference']['state_dict_path']
assert Path(net_params['inference']['state_dict_path']).is_file(
), f"Could not locate {net_params['inference']['state_dict_path']}"
checkpoint = load_checkpoint(state_dict_path)
elif train_state_dict_path is not None:
assert Path(train_state_dict_path).is_file(
), f'Could not locate checkpoint at {train_state_dict_path}'
checkpoint = load_checkpoint(train_state_dict_path)
else:
checkpoint = None
return model, checkpoint, model_name
def main(data_path, batch_size, num_epochs, start_epoch, learning_rate,
momentum, print_freq, run, resume, data_set):
"""
Args:
data_path:
batch_size:
num_epochs:
Returns:
"""
since = time.time()
# get model
model = unet.UNetSmall()
if torch.cuda.is_available():
model = model.cuda()
# set up binary cross entropy and dice loss
criterion = metrics.BCEDiceLoss()
# optimizer
# optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, nesterov=True)
# optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-5)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# decay LR
lr_scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=18,
gamma=0.1)
# starting params
best_loss = 999
# optionally resume from a checkpoint
if resume:
if os.path.isfile(resume):
print("=> loading checkpoint '{}'".format(resume))
checkpoint = torch.load(resume)
if checkpoint['epoch'] > start_epoch:
start_epoch = checkpoint['epoch']
best_loss = checkpoint['best_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# get data
mass_dataset_train = data_utils.MassRoadBuildingDataset(
data_path,
data_set,
'train',
transform=transforms.Compose([aug.ToTensorTarget()]))
mass_dataset_val = data_utils.MassRoadBuildingDataset(
data_path,
data_set,
'valid',
transform=transforms.Compose([aug.ToTensorTarget()]))
# creating loaders
train_dataloader = DataLoader(mass_dataset_train,
batch_size=batch_size,
num_workers=2,
shuffle=True)
val_dataloader = DataLoader(mass_dataset_val,
batch_size=3,
num_workers=2,
shuffle=False)
# loggers
train_logger = logger.Logger('../logs/run_{}/training'.format(str(run)),
print_freq)
val_logger = logger.Logger('../logs/run_{}/validation'.format(str(run)),
print_freq)
for epoch in range(start_epoch, num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
# step the learning rate scheduler
lr_scheduler.step()
# run training and validation
train_metrics = train(train_dataloader, model, criterion, optimizer,
lr_scheduler, train_logger, epoch)
valid_metrics = validation(val_dataloader, model, criterion,
val_logger, epoch)
# store best loss and save a model checkpoint
is_best = valid_metrics['valid_loss'] < best_loss
best_loss = min(valid_metrics['valid_loss'], best_loss)
save_checkpoint(
{
'epoch': epoch,
'arch': 'UNetSmall',
'state_dict': model.state_dict(),
'best_loss': best_loss,
'optimizer': optimizer.state_dict()
}, is_best)
cur_elapsed = time.time() - since
print('Current elapsed time {:.0f}m {:.0f}s'.format(
cur_elapsed // 60, cur_elapsed % 60))
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
def net(net_params, num_channels, inference=False):
"""Define the neural net"""
model_name = net_params['global']['model_name'].lower()
num_bands = int(net_params['global']['number_of_bands'])
msg = f'Number of bands specified incompatible with this model. Requires 3 band data.'
train_state_dict_path = get_key_def('state_dict_path',
net_params['training'], None)
pretrained = get_key_def('pretrained', net_params['training'],
True) if not inference else False
dropout = get_key_def('dropout', net_params['training'], False)
dropout_prob = get_key_def('dropout_prob', net_params['training'], 0.5)
# TODO: find a way to maybe implement it in classification one day
if 'concatenate_depth' in net_params['global']:
# Read the concatenation point
conc_point = net_params['global']['concatenate_depth']
if model_name == 'unetsmall':
model = unet.UNetSmall(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'unet':
model = unet.UNet(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'ternausnet':
assert num_bands == 3, msg
model = TernausNet.ternausnet(num_channels)
elif model_name == 'checkpointed_unet':
model = checkpointed_unet.UNetSmall(num_channels, num_bands, dropout,
dropout_prob)
elif model_name == 'inception':
model = inception.Inception3(num_channels, num_bands)
elif model_name == 'fcn_resnet101':
assert num_bands == 3, msg
model = models.segmentation.fcn_resnet101(pretrained=pretrained,
progress=True,
num_classes=num_channels,
aux_loss=None)
elif model_name == 'deeplabv3_resnet101':
assert (num_bands == 3 or num_bands == 4), msg
if num_bands == 3:
print('Finetuning pretrained deeplabv3 with 3 bands')
model = models.segmentation.deeplabv3_resnet101(
pretrained=pretrained, progress=True)
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module(
'4',
nn.Conv2d(classifier[-1].in_channels,
num_channels,
kernel_size=(1, 1)))
elif num_bands == 4:
print('Finetuning pretrained deeplabv3 with 4 bands')
print('Testing with 4 bands, concatenating at {}.'.format(
conc_point))
model = models.segmentation.deeplabv3_resnet101(
pretrained=pretrained, progress=True)
if conc_point == 'baseline':
conv1 = model.backbone._modules['conv1'].weight.detach().numpy(
)
depth = np.expand_dims(
conv1[:, 1,
...], axis=1) # reuse green weights for infrared.
conv1 = np.append(conv1, depth, axis=1)
conv1 = torch.from_numpy(conv1).float()
model.backbone._modules['conv1'].weight = nn.Parameter(
conv1, requires_grad=True)
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module(
'4',
nn.Conv2d(classifier[-1].in_channels,
num_channels,
kernel_size=(1, 1)))
else:
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module(
'4',
nn.Conv2d(classifier[-1].in_channels,
num_channels,
kernel_size=(1, 1)))
###################
#conv1 = model.backbone._modules['conv1'].weight.detach().numpy()
#depth = np.random.uniform(low=-1, high=1, size=(64, 1, 7, 7))
#conv1 = np.append(conv1, depth, axis=1)
#conv1 = torch.from_numpy(conv1).float()
#model.backbone._modules['conv1'].weight = nn.Parameter(conv1, requires_grad=True)
###################
model = LayersEnsemble(model, conc_point=conc_point)
elif model_name in lm_smp.keys():
lsmp = lm_smp[model_name]
# TODO: add possibility of our own weights
lsmp['params'][
'encoder_weights'] = "imagenet" if 'pretrained' in model_name.split(
"_") else None
lsmp['params']['in_channels'] = num_bands
lsmp['params']['classes'] = num_channels
lsmp['params']['activation'] = None
model = lsmp['fct'](**lsmp['params'])
else:
raise ValueError(
f'The model name {model_name} in the config.yaml is not defined.')
coordconv_convert = get_key_def('coordconv_convert', net_params['global'],
False)
if coordconv_convert:
centered = get_key_def('coordconv_centered', net_params['global'],
True)
normalized = get_key_def('coordconv_normalized', net_params['global'],
True)
noise = get_key_def('coordconv_noise', net_params['global'], None)
radius_channel = get_key_def('coordconv_radius_channel',
net_params['global'], False)
scale = get_key_def('coordconv_scale', net_params['global'], 1.0)
# note: this operation will not attempt to preserve already-loaded model parameters!
model = coordconv.swap_coordconv_layers(model,
centered=centered,
normalized=normalized,
noise=noise,
radius_channel=radius_channel,
scale=scale)
if inference:
state_dict_path = net_params['inference']['state_dict_path']
assert Path(net_params['inference']['state_dict_path']).is_file(
), f"Could not locate {net_params['inference']['state_dict_path']}"
checkpoint = load_checkpoint(state_dict_path)
elif train_state_dict_path is not None:
assert Path(train_state_dict_path).is_file(
), f'Could not locate checkpoint at {train_state_dict_path}'
checkpoint = load_checkpoint(train_state_dict_path)
else:
checkpoint = None
return model, checkpoint, model_name
def net(net_params, num_channels, inference=False):
"""Define the neural net"""
model_name = net_params['global']['model_name'].lower()
num_bands = int(net_params['global']['number_of_bands'])
msg = f'Number of bands specified incompatible with this model. Requires 3 band data.'
train_state_dict_path = get_key_def('state_dict_path', net_params['training'], None)
pretrained = get_key_def('pretrained', net_params['training'], True) if not inference else False
dropout = get_key_def('dropout', net_params['training'], False)
dropout_prob = get_key_def('dropout_prob', net_params['training'], 0.5)
if model_name == 'unetsmall':
model = unet.UNetSmall(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'unet':
model = unet.UNet(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'ternausnet':
assert num_bands == 3, msg
model = TernausNet.ternausnet(num_channels)
elif model_name == 'checkpointed_unet':
model = checkpointed_unet.UNetSmall(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'inception':
model = inception.Inception3(num_channels, num_bands)
elif model_name == 'fcn_resnet101':
assert num_bands == 3, msg
model = models.segmentation.fcn_resnet101(pretrained=False, progress=True, num_classes=num_channels,
aux_loss=None)
elif model_name == 'deeplabv3_resnet101':
try:
model = models.segmentation.deeplabv3_resnet101(pretrained=False, progress=True, in_channels=num_bands,
num_classes=num_channels, aux_loss=None)
except:
assert num_bands==3, 'Edit torchvision scripts segmentation.py and resnet.py to build deeplabv3_resnet ' \
'with more or less than 3 bands'
model = models.segmentation.deeplabv3_resnet101(pretrained=False, progress=True,
num_classes=num_channels, aux_loss=None)
else:
raise ValueError(f'The model name {model_name} in the config.yaml is not defined.')
coordconv_convert = get_key_def('coordconv_convert', net_params['global'], False)
if coordconv_convert:
centered = get_key_def('coordconv_centered', net_params['global'], True)
normalized = get_key_def('coordconv_normalized', net_params['global'], True)
noise = get_key_def('coordconv_noise', net_params['global'], None)
radius_channel = get_key_def('coordconv_radius_channel', net_params['global'], False)
scale = get_key_def('coordconv_scale', net_params['global'], 1.0)
# note: this operation will not attempt to preserve already-loaded model parameters!
model = coordconv.swap_coordconv_layers(model, centered=centered, normalized=normalized, noise=noise,
radius_channel=radius_channel, scale=scale)
if inference:
state_dict_path = net_params['inference']['state_dict_path']
assert Path(net_params['inference']['state_dict_path']).is_file(), f"Could not locate {net_params['inference']['state_dict_path']}"
checkpoint = load_checkpoint(state_dict_path)
elif train_state_dict_path is not None:
assert Path(train_state_dict_path).is_file(), f'Could not locate checkpoint at {train_state_dict_path}'
checkpoint = load_checkpoint(train_state_dict_path)
elif pretrained and (model_name == ('deeplabv3_resnet101' or 'fcn_resnet101')):
print(f'Retrieving coco checkpoint for {model_name}...\n')
if model_name == 'deeplabv3_resnet101': # default to pretrained on coco (21 classes)
coco_model = models.segmentation.deeplabv3_resnet101(pretrained=True, progress=True, num_classes=21, aux_loss=None)
else:
coco_model = models.segmentation.fcn_resnet101(pretrained=True, progress=True, num_classes=21, aux_loss=None)
checkpoint = coco_model.state_dict()
# Place entire state_dict inside 'model' key for compatibility with the rest of GDL workflow
temp_checkpoint = {}
temp_checkpoint['model'] = {k: v for k, v in checkpoint.items()}
del coco_model, checkpoint
checkpoint = temp_checkpoint
elif pretrained:
warnings.warn(f'No pretrained checkpoint found for {model_name}.')
checkpoint = None
else:
checkpoint = None
return model, checkpoint, model_name
def net(net_params, num_channels, inference=False):
"""Define the neural net"""
model_name = net_params['global']['model_name'].lower()
num_bands = int(net_params['global']['number_of_bands'])
msg = f'Number of bands specified incompatible with this model. Requires 3 band data.'
train_state_dict_path = get_key_def('state_dict_path', net_params['training'], None)
pretrained = get_key_def('pretrained', net_params['training'], True) if not inference else False
dropout = get_key_def('dropout', net_params['training'], False)
dropout_prob = get_key_def('dropout_prob', net_params['training'], 0.5)
dontcare_val = get_key_def("ignore_index", net_params["training"], -1)
num_devices = net_params['global']['num_gpus']
if dontcare_val == 0:
warnings.warn("The 'dontcare' value (or 'ignore_index') used in the loss function cannot be zero;"
" all valid class indices should be consecutive, and start at 0. The 'dontcare' value"
" will be remapped to -1 while loading the dataset, and inside the config from now on.")
net_params["training"]["ignore_index"] = -1
# TODO: find a way to maybe implement it in classification one day
if 'concatenate_depth' in net_params['global']:
# Read the concatenation point
conc_point = net_params['global']['concatenate_depth']
if model_name == 'unetsmall':
model = unet.UNetSmall(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'unet':
model = unet.UNet(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'ternausnet':
assert num_bands == 3, msg
model = TernausNet.ternausnet(num_channels)
elif model_name == 'checkpointed_unet':
model = checkpointed_unet.UNetSmall(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'inception':
model = inception.Inception3(num_channels, num_bands)
elif model_name == 'fcn_resnet101':
assert num_bands == 3, msg
model = models.segmentation.fcn_resnet101(pretrained=False, progress=True, num_classes=num_channels,
aux_loss=None)
elif model_name == 'deeplabv3_resnet101':
assert (num_bands == 3 or num_bands == 4), msg
if num_bands == 3:
print('Finetuning pretrained deeplabv3 with 3 bands')
model = models.segmentation.deeplabv3_resnet101(pretrained=pretrained, progress=True)
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module('4', nn.Conv2d(classifier[-1].in_channels, num_channels, kernel_size=(1, 1)))
elif num_bands == 4:
print('Finetuning pretrained deeplabv3 with 4 bands')
print('Testing with 4 bands, concatenating at {}.'.format(conc_point))
model = models.segmentation.deeplabv3_resnet101(pretrained=pretrained, progress=True)
if conc_point=='baseline':
conv1 = model.backbone._modules['conv1'].weight.detach().numpy()
depth = np.expand_dims(conv1[:, 1, ...], axis=1) # reuse green weights for infrared.
conv1 = np.append(conv1, depth, axis=1)
conv1 = torch.from_numpy(conv1).float()
model.backbone._modules['conv1'].weight = nn.Parameter(conv1, requires_grad=True)
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module(
'4', nn.Conv2d(classifier[-1].in_channels, num_channels, kernel_size=(1, 1))
)
else:
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module(
'4', nn.Conv2d(classifier[-1].in_channels, num_channels, kernel_size=(1, 1))
)
###################
#conv1 = model.backbone._modules['conv1'].weight.detach().numpy()
#depth = np.random.uniform(low=-1, high=1, size=(64, 1, 7, 7))
#conv1 = np.append(conv1, depth, axis=1)
#conv1 = torch.from_numpy(conv1).float()
#model.backbone._modules['conv1'].weight = nn.Parameter(conv1, requires_grad=True)
###################
model = LayersEnsemble(model, conc_point=conc_point)
elif model_name in lm_smp.keys():
lsmp = lm_smp[model_name]
# TODO: add possibility of our own weights
lsmp['params']['encoder_weights'] = "imagenet" if 'pretrained' in model_name.split("_") else None
lsmp['params']['in_channels'] = num_bands
lsmp['params']['classes'] = num_channels
lsmp['params']['activation'] = None
model = lsmp['fct'](**lsmp['params'])
else:
raise ValueError(f'The model name {model_name} in the config.yaml is not defined.')
coordconv_convert = get_key_def('coordconv_convert', net_params['global'], False)
if coordconv_convert:
centered = get_key_def('coordconv_centered', net_params['global'], True)
normalized = get_key_def('coordconv_normalized', net_params['global'], True)
noise = get_key_def('coordconv_noise', net_params['global'], None)
radius_channel = get_key_def('coordconv_radius_channel', net_params['global'], False)
scale = get_key_def('coordconv_scale', net_params['global'], 1.0)
# note: this operation will not attempt to preserve already-loaded model parameters!
model = coordconv.swap_coordconv_layers(model, centered=centered, normalized=normalized, noise=noise,
radius_channel=radius_channel, scale=scale)
if inference:
state_dict_path = net_params['inference']['state_dict_path']
assert Path(net_params['inference']['state_dict_path']).is_file(), f"Could not locate {net_params['inference']['state_dict_path']}"
checkpoint = load_checkpoint(state_dict_path)
return model, checkpoint, model_name
else:
if train_state_dict_path is not None:
assert Path(train_state_dict_path).is_file(), f'Could not locate checkpoint at {train_state_dict_path}'
checkpoint = load_checkpoint(train_state_dict_path)
else:
checkpoint = None
assert num_devices is not None and num_devices >= 0, "missing mandatory num gpus parameter"
# list of GPU devices that are available and unused. If no GPUs, returns empty list
lst_device_ids = get_device_ids(num_devices) if torch.cuda.is_available() else []
num_devices = len(lst_device_ids) if lst_device_ids else 0
device = torch.device(f'cuda:{lst_device_ids[0]}' if torch.cuda.is_available() and lst_device_ids else 'cpu')
print(f"Number of cuda devices requested: {net_params['global']['num_gpus']}. Cuda devices available: {lst_device_ids}\n")
if num_devices == 1:
print(f"Using Cuda device {lst_device_ids[0]}\n")
elif num_devices > 1:
print(f"Using data parallel on devices: {str(lst_device_ids)[1:-1]}. Main device: {lst_device_ids[0]}\n") # TODO: why are we showing indices [1:-1] for lst_device_ids?
try: # For HPC when device 0 not available. Error: Invalid device id (in torch/cuda/__init__.py).
model = nn.DataParallel(model,
device_ids=lst_device_ids) # DataParallel adds prefix 'module.' to state_dict keys
except AssertionError:
warnings.warn(f"Unable to use devices {lst_device_ids}. Trying devices {list(range(len(lst_device_ids)))}")
device = torch.device('cuda:0')
lst_device_ids = range(len(lst_device_ids))
model = nn.DataParallel(model,
device_ids=lst_device_ids) # DataParallel adds prefix 'module.' to state_dict keys
else:
warnings.warn(f"No Cuda device available. This process will only run on CPU\n")
tqdm.write(f'Setting model, criterion, optimizer and learning rate scheduler...\n')
try: # For HPC when device 0 not available. Error: Cuda invalid device ordinal.
model.to(device)
except RuntimeError:
warnings.warn(f"Unable to use device. Trying device 0...\n")
device = torch.device(f'cuda:0' if torch.cuda.is_available() and lst_device_ids else 'cpu')
model.to(device)
model, criterion, optimizer, lr_scheduler = set_hyperparameters(net_params, num_channels, model, checkpoint, dontcare_val)
criterion = criterion.to(device)
return model, model_name, criterion, optimizer, lr_scheduler
def segmenter(run, data_num, in_folder, out_csv, start, stop):
#outpath_csv, outcsv = files2csv(in_folder, out_csv)
outpath_csv = "../../assets/output_data/csv_files/" + out_csv
model = unet.UNetSmall()
model.eval() # setting to eval
optimizer = optim.Adam(model.parameters(), lr=0.001)
if torch.cuda.is_available():
print("Cuda is available.")
model = model.cuda()
checkpoint = torch.load('../checkpoints/model_best.pth.tar')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
mass_dataset_test = data_utils.MassRoadBuildingDataset(
run,
data_num,
outpath_csv,
start,
stop,
in_folder,
'test',
transform=transforms.Compose([aug.ToTensorTarget()]))
test_dataloader = DataLoader(mass_dataset_test,
batch_size=1,
num_workers=3,
shuffle=False)
seg_res = []
sat_imgs = []
segres_unrolled = []
sat_unrolled = []
counter = 0
with torch.no_grad():
for idx, data in enumerate(
tqdm(test_dataloader, desc='test' + " run= " + str(run))):
# get the inputs and wrap in Variable
if torch.cuda.is_available():
inputs = Variable(data['sat_img'].cuda(), volatile=True)
labels = Variable(data['sat_img'].cuda(), volatile=True)
else:
inputs = Variable(data['sat_img'], volatile=True)
labels = Variable(data['sat_img'], volatile=True)
optimizer.zero_grad()
outputs = model(inputs)
outputs = torch.nn.functional.sigmoid(outputs)
seg_res.append(outputs)
sat_imgs.append(inputs)
counter += 1
# unroll segmentation output
segres_unrolled = [
np.squeeze((seg_res[i].data).cpu().numpy())
for i in range(len(seg_res))
]
# unroll sat_imgs
sat_unrolled = [
np.squeeze((sat_imgs[i].data).cpu().numpy())
for i in range(len(sat_imgs))
]
return segres_unrolled, sat_unrolled
def net(net_params, inference=False):
"""Define the neural net"""
model_name = net_params['global']['model_name'].lower()
num_classes = net_params['global']['num_classes']
msg = f'Number of bands specified incompatible with this model. Requires 3 band data.'
state_dict_path = ''
if model_name == 'unetsmall':
model = unet.UNetSmall(
num_classes, net_params['global']['number_of_bands'],
net_params['models']['unetsmall']['dropout'],
net_params['models']['unetsmall']['probability'])
if net_params['models']['unetsmall']['pretrained']:
state_dict_path = net_params['models']['unetsmall']['pretrained']
elif model_name == 'unet':
model = unet.UNet(num_classes, net_params['global']['number_of_bands'],
net_params['models']['unet']['dropout'],
net_params['models']['unet']['probability'])
if net_params['models']['unet']['pretrained']:
state_dict_path = net_params['models']['unet']['pretrained']
elif model_name == 'ternausnet':
model = TernausNet.ternausnet(
num_classes, net_params['models']['ternausnet']['pretrained'])
elif model_name == 'checkpointed_unet':
model = checkpointed_unet.UNetSmall(
num_classes, net_params['global']['number_of_bands'],
net_params['models']['unetsmall']['dropout'],
net_params['models']['unetsmall']['probability'])
if net_params['models']['unetsmall']['pretrained']:
state_dict_path = net_params['models']['unetsmall']['pretrained']
elif model_name == 'inception':
model = inception.Inception3(num_classes,
net_params['global']['number_of_bands'])
if net_params['models']['inception']['pretrained']:
state_dict_path = net_params['models']['inception']['pretrained']
elif model_name == 'fcn_resnet101':
assert net_params['global']['number_of_bands'], msg
coco_model = models.segmentation.fcn_resnet101(pretrained=True,
progress=True,
num_classes=21,
aux_loss=None)
model = models.segmentation.fcn_resnet101(pretrained=False,
progress=True,
num_classes=num_classes,
aux_loss=None)
chopped_dict = chop_layer(coco_model.state_dict(),
layer_name='classifier.4')
del coco_model
model.load_state_dict(chopped_dict,
strict=False) # load the new state dict
if net_params['models']['fcn_resnet101']['pretrained']:
state_dict_path = net_params['models']['fcn_resnet101'][
'pretrained']
elif model_name == 'deeplabv3_resnet101':
assert net_params['global']['number_of_bands'], msg
# pretrained on coco (21 classes)
coco_model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True,
num_classes=21,
aux_loss=None)
model = models.segmentation.deeplabv3_resnet101(
pretrained=False,
progress=True,
num_classes=num_classes,
aux_loss=None)
chopped_dict = chop_layer(coco_model.state_dict(),
layer_name='classifier.4')
del coco_model
# load the new state dict
model.load_state_dict(
chopped_dict, strict=False
) # When strict=False, allows to load only the variables that
# are identical between the two models irrespective of whether one is subset/superset of the other.
if net_params['models']['deeplabv3_resnet101']['pretrained']:
state_dict_path = net_params['models']['deeplabv3_resnet101'][
'pretrained']
else:
raise ValueError('The model name in the config.yaml is not defined.')
if inference:
state_dict_path = net_params['inference']['state_dict_path']
return model, state_dict_path, model_name
def net(net_params, inference=False):
"""Define the neural net"""
model_name = net_params['global']['model_name'].lower()
num_classes = net_params['global']['num_classes']
if num_classes == 1:
warnings.warn(
"config specified that number of classes is 1, but model will be instantiated"
" with a minimum of two regardless (will assume that 'background' exists)"
)
num_classes = 2
msg = f'Number of bands specified incompatible with this model. Requires 3 band data.'
state_dict_path = ''
if model_name == 'unetsmall':
model = unet.UNetSmall(num_classes,
net_params['global']['number_of_bands'],
net_params['training']['dropout'],
net_params['training']['dropout_prob'])
elif model_name == 'unet':
model = unet.UNet(num_classes, net_params['global']['number_of_bands'],
net_params['training']['dropout'],
net_params['training']['dropout_prob'])
elif model_name == 'ternausnet':
assert net_params['global']['number_of_bands'] == 3, msg
model = TernausNet.ternausnet(num_classes)
elif model_name == 'checkpointed_unet':
model = checkpointed_unet.UNetSmall(
num_classes, net_params['global']['number_of_bands'],
net_params['training']['dropout'],
net_params['training']['dropout_prob'])
elif model_name == 'inception':
model = inception.Inception3(num_classes,
net_params['global']['number_of_bands'])
elif model_name == 'fcn_resnet101':
assert net_params['global']['number_of_bands'] == 3, msg
coco_model = models.segmentation.fcn_resnet101(pretrained=True,
progress=True,
num_classes=21,
aux_loss=None)
model = models.segmentation.fcn_resnet101(pretrained=False,
progress=True,
num_classes=num_classes,
aux_loss=None)
chopped_dict = chop_layer(coco_model.state_dict(),
layer_names=['classifier.4'])
del coco_model
# load the new state dict
# When strict=False, allows to load only the variables that are identical between the two models irrespective of
# whether one is subset/superset of the other.
model.load_state_dict(chopped_dict, strict=False)
elif model_name == 'deeplabv3_resnet101':
assert net_params['global']['number_of_bands'] == 3, msg
# pretrained on coco (21 classes)
coco_model = models.segmentation.deeplabv3_resnet101(pretrained=True,
progress=True,
num_classes=21,
aux_loss=None)
model = models.segmentation.deeplabv3_resnet101(
pretrained=False,
progress=True,
num_classes=num_classes,
aux_loss=None)
chopped_dict = chop_layer(coco_model.state_dict(),
layer_names=['classifier.4'])
del coco_model
model.load_state_dict(chopped_dict, strict=False)
else:
raise ValueError(
f'The model name {model_name} in the config.yaml is not defined.')
coordconv_convert = get_key_def('coordconv_convert', net_params['global'],
False)
if coordconv_convert:
centered = get_key_def('coordconv_centered', net_params['global'],
True)
normalized = get_key_def('coordconv_normalized', net_params['global'],
True)
noise = get_key_def('coordconv_noise', net_params['global'], None)
radius_channel = get_key_def('coordconv_radius_channel',
net_params['global'], False)
scale = get_key_def('coordconv_scale', net_params['global'], 1.0)
# note: this operation will not attempt to preserve already-loaded model parameters!
model = coordconv.swap_coordconv_layers(model,
centered=centered,
normalized=normalized,
noise=noise,
radius_channel=radius_channel,
scale=scale)
if net_params['training']['state_dict_path']:
state_dict_path = net_params['training']['state_dict_path']
checkpoint = load_checkpoint(state_dict_path)
elif inference:
state_dict_path = net_params['inference']['state_dict_path']
checkpoint = load_checkpoint(state_dict_path)
else:
checkpoint = None
return model, checkpoint, model_name
def net(model_name: str,
num_bands: int,
num_channels: int,
dontcare_val: int,
num_devices: int,
train_state_dict_path: str = None,
pretrained: bool = True,
dropout_prob: float = False,
loss_fn: str = None,
optimizer: str = None,
class_weights: Sequence = None,
net_params=None,
conc_point: str = None,
coordconv_params=None,
inference_state_dict: str = None):
"""Define the neural net"""
msg = f'Number of bands specified incompatible with this model. Requires 3 band data.'
pretrained = False if train_state_dict_path or inference_state_dict else pretrained
dropout = True if dropout_prob else False
model = None
if model_name == 'unetsmall':
model = unet.UNetSmall(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'unet':
model = unet.UNet(num_channels, num_bands, dropout, dropout_prob)
elif model_name == 'ternausnet':
if not num_bands == 3:
raise NotImplementedError(msg)
model = TernausNet.ternausnet(num_channels)
elif model_name == 'checkpointed_unet':
model = checkpointed_unet.UNetSmall(num_channels, num_bands, dropout,
dropout_prob)
elif model_name == 'inception':
model = inception.Inception3(num_channels, num_bands)
elif model_name == 'fcn_resnet101':
if not num_bands == 3:
raise NotImplementedError(msg)
model = models.segmentation.fcn_resnet101(pretrained=False,
progress=True,
num_classes=num_channels,
aux_loss=None)
elif model_name == 'deeplabv3_resnet101':
if not (num_bands == 3 or num_bands == 4):
raise NotImplementedError(msg)
if num_bands == 3:
model = models.segmentation.deeplabv3_resnet101(
pretrained=pretrained, progress=True)
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module(
'4',
nn.Conv2d(classifier[-1].in_channels,
num_channels,
kernel_size=(1, 1)))
elif num_bands == 4:
model = models.segmentation.deeplabv3_resnet101(
pretrained=pretrained, progress=True)
if conc_point == 'baseline':
logging.info(
'Testing with 4 bands, concatenating at {}.'.format(
conc_point))
conv1 = model.backbone._modules['conv1'].weight.detach().numpy(
)
depth = np.expand_dims(
conv1[:, 1,
...], axis=1) # reuse green weights for infrared.
conv1 = np.append(conv1, depth, axis=1)
conv1 = torch.from_numpy(conv1).float()
model.backbone._modules['conv1'].weight = nn.Parameter(
conv1, requires_grad=True)
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module(
'4',
nn.Conv2d(classifier[-1].in_channels,
num_channels,
kernel_size=(1, 1)))
else:
classifier = list(model.classifier.children())
model.classifier = nn.Sequential(*classifier[:-1])
model.classifier.add_module(
'4',
nn.Conv2d(classifier[-1].in_channels,
num_channels,
kernel_size=(1, 1)))
###################
# conv1 = model.backbone._modules['conv1'].weight.detach().numpy()
# depth = np.random.uniform(low=-1, high=1, size=(64, 1, 7, 7))
# conv1 = np.append(conv1, depth, axis=1)
# conv1 = torch.from_numpy(conv1).float()
# model.backbone._modules['conv1'].weight = nn.Parameter(conv1, requires_grad=True)
###################
conc_point = 'conv1' if not conc_point else conc_point
model = LayersEnsemble(model, conc_point=conc_point)
logging.info(
f'Finetuning pretrained deeplabv3 with {num_bands} input channels (imagery bands). '
f'Concatenation point: "{conc_point}"')
elif model_name in lm_smp.keys():
lsmp = lm_smp[model_name]
# TODO: add possibility of our own weights
lsmp['params'][
'encoder_weights'] = "imagenet" if 'pretrained' in model_name.split(
"_") else None
lsmp['params']['in_channels'] = num_bands
lsmp['params']['classes'] = num_channels
lsmp['params']['activation'] = None
model = lsmp['fct'](**lsmp['params'])
else:
raise ValueError(
f'The model name {model_name} in the config.yaml is not defined.')
coordconv_convert = get_key_def('coordconv_convert', coordconv_params,
False)
if coordconv_convert:
centered = get_key_def('coordconv_centered', coordconv_params, True)
normalized = get_key_def('coordconv_normalized', coordconv_params,
True)
noise = get_key_def('coordconv_noise', coordconv_params, None)
radius_channel = get_key_def('coordconv_radius_channel',
coordconv_params, False)
scale = get_key_def('coordconv_scale', coordconv_params, 1.0)
# note: this operation will not attempt to preserve already-loaded model parameters!
model = coordconv.swap_coordconv_layers(model,
centered=centered,
normalized=normalized,
noise=noise,
radius_channel=radius_channel,
scale=scale)
if inference_state_dict:
state_dict_path = inference_state_dict
checkpoint = load_checkpoint(state_dict_path)
return model, checkpoint, model_name
else:
if train_state_dict_path is not None:
checkpoint = load_checkpoint(train_state_dict_path)
else:
checkpoint = None
# list of GPU devices that are available and unused. If no GPUs, returns empty list
gpu_devices_dict = get_device_ids(num_devices)
num_devices = len(gpu_devices_dict.keys())
logging.info(
f"Number of cuda devices requested: {num_devices}. "
f"Cuda devices available: {list(gpu_devices_dict.keys())}\n")
if num_devices == 1:
logging.info(
f"Using Cuda device 'cuda:{list(gpu_devices_dict.keys())[0]}'")
elif num_devices > 1:
logging.info(
f"Using data parallel on devices: {list(gpu_devices_dict.keys())[1:]}. "
f"Main device: 'cuda:{list(gpu_devices_dict.keys())[0]}'")
try: # For HPC when device 0 not available. Error: Invalid device id (in torch/cuda/__init__.py).
# DataParallel adds prefix 'module.' to state_dict keys
model = nn.DataParallel(model,
device_ids=list(
gpu_devices_dict.keys()))
except AssertionError:
logging.warning(
f"Unable to use devices with ids {gpu_devices_dict.keys()}"
f"Trying devices with ids {list(range(len(gpu_devices_dict.keys())))}"
)
model = nn.DataParallel(
model,
device_ids=list(range(len(gpu_devices_dict.keys()))))
else:
logging.warning(
f"No Cuda device available. This process will only run on CPU\n"
)
logging.info(
f'Setting model, criterion, optimizer and learning rate scheduler...\n'
)
device = torch.device(
f'cuda:{list(range(len(gpu_devices_dict.keys())))[0]}'
if gpu_devices_dict else 'cpu')
try: # For HPC when device 0 not available. Error: Cuda invalid device ordinal.
model.to(device)
except AssertionError:
logging.exception(f"Unable to use device. Trying device 0...\n")
device = torch.device(f'cuda' if gpu_devices_dict else 'cpu')
model.to(device)
model, criterion, optimizer, lr_scheduler = set_hyperparameters(
params=net_params,
num_classes=num_channels,
model=model,
checkpoint=checkpoint,
dontcare_val=dontcare_val,
loss_fn=loss_fn,
optimizer=optimizer,
class_weights=class_weights,
inference=inference_state_dict)
criterion = criterion.to(device)
return model, model_name, criterion, optimizer, lr_scheduler, device, gpu_devices_dict
