def train_on_tiles(data_dir,
model_id,
max_epochs=25,
learning_rate=1e-3,
batch_size=32,
threshold=0.5,
n_classes=None,
class_weights=None,
batch_norm=True,
augmentation=False,
input_channels=None,
depth=4,
wf=4,
padding=True,
scheduler=None):
print('Building U-Net model and training')
exports_dir = './figures/{}/'.format(model_id)
os.makedirs(exports_dir, exist_ok=True)
# Define augmentation transforms
aug = iaa.Sequential([
# iaa.MultiplyHueAndSaturation((0.5, 1.5)),
# iaa.AddToHueAndSaturation((-45, 45)),
# iaa.Grayscale((0.0, 1.0)),
# iaa.AllChannelsHistogramEqualization(),
# iaa.GammaContrast((0.0, 1.75), per_channel=True),
# iaa.LinearContrast((0.0, 1.75), per_channel=True),
iaa.Crop(px=(0, 32)), # randomly crop between 0 and 32 pixels
iaa.Fliplr(0.50), # horizontally flip 50% of the images
iaa.Flipud(0.50), # horizontally flip 50% of the images
iaa.Rot90([0, 1, 2, 3]), # apply any of the 90 deg rotations
# iaa.PerspectiveTransform((0.025, 0.1)), # randomly scale between 0.025 and 0.1
# iaa.Affine(scale=(0.5, 1.5), translate_percent=(0, 0.25), # random affine transforms with symmetric padding
# rotate=(0, 360), shear=(0, 360), mode='symmetric'),
])
if class_weights is None:
class_weights = [1, 1, 1]
if augmentation:
tfms = ImgAugTransform(aug)
else:
tfms = None
datasets = {}
datasets['train'] = OrthoData(os.path.join(data_dir, 'train_tiles.h5'),
transform=None,
channels=input_channels)
datasets['val'] = OrthoData(os.path.join(data_dir, 'test_tiles.h5'),
transform=None,
channels=input_channels)
if input_channels is None:
n_channels = datasets['train'].input_shape[1]
else:
n_channels = len(input_channels)
if n_classes is None:
n_classes = datasets['train'].n_classes
print('Input channels:', input_channels)
# DataLoader parameters
dataloader_params = {
'batch_size': batch_size,
'num_workers': 0,
'shuffle': True
}
# Make the DataLoaders
data_loaders = {
name: DataLoader(datasets[name], **dataloader_params)
for name in ['train', 'val']
}
# Define model
print('CUDA:', torch.cuda.is_available())
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
unet_params = {
'n_classes': n_classes,
'in_channels': n_channels,
'depth': depth,
'padding': padding,
'wf': wf,
'up_mode': 'upconv',
'batch_norm': batch_norm
}
model = UNet(**unet_params).to(device)
n_parameters = count_parameters(model)
print(f'Number of model parameters: {n_parameters}')
# Define optimizer and loss function
optim = torch.optim.Adam(model.parameters(), lr=learning_rate)
class_weights = torch.Tensor(class_weights).to(device)
criterion = nn.CrossEntropyLoss(class_weights)
if scheduler:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim,
verbose=True,
patience=2)
scheduler = GradualWarmupScheduler(
optim,
multiplier=10,
total_epoch=10,
after_scheduler=scheduler,
)
# Train
model, history = train_model(model,
data_loaders,
max_epochs,
criterion,
optim,
scheduler=scheduler,
return_history=True,
transform=tfms)
os.makedirs('./saved_weights/', exist_ok=True)
model_filename = '{}.pt'.format(model_id)
torch.save(model.state_dict(), './saved_weights/{}'.format(model_filename))
print('Saved model as: ./saved_weights/{}'.format(model_filename))
n_examples = len(datasets['val']) # //10
sampler = RandomSampler(datasets['val'],
replacement=True,
num_samples=n_examples)
inputs = DataLoader(datasets['val'], sampler=sampler)
print('Predicting over the validation dataset')
inputs, preds, probs, targets = predict(model,
data_loaders['val'],
th=threshold)
plot_results(inputs=inputs.astype(np.uint8),
classes=np.expand_dims(targets, 1),
prob=probs,
n_plot=5,
save_path=exports_dir + 'example_tiles.png')
learning_curve(history=history,
name='loss',
outfile=exports_dir + 'loss_learning_curves.png')
compare_learning_curve(history=history,
name='accuracy',
outfile=exports_dir +
'accuracy_learning_curves.png')
