dataset = BasicDataset(unet_type, dir_img, dir_mask, img_scale)

n_val = int(len(dataset) * val_percent)

n_train = len(dataset) - n_val

train, val = random_split(dataset, [n_train, n_val])

train_loader = DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)

val_loader = DataLoader(val, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True)

writer = SummaryWriter(comment=f'LR_{lr}_BS_{batch_size}_SCALE_{img_scale}')

global_step = 0

logging.info(f'''Starting training:

UNet type: {unet_type}

Epochs: {epochs}

Batch size: {batch_size}

Learning rate: {lr}

Dataset size: {len(dataset)}

Training size: {n_train}

Validation size: {n_val}

Checkpoints: {save_cp}

Device: {device.type}

Images scaling: {img_scale}''')

optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8)

model, optimizer = amp.initialize(net, optimizer, opt_level="O1")

# Scheduler https://arxiv.org/pdf/1812.01187.pdf

lf = lambda x: (((1 math.cos(x * math.pi / epochs)) / 2) ** 1.0) * 0.95 0.05 #cosine

scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

scheduler.last_epoch = global_step

if net.n_classes > 1:

criterion = nn.CrossEntropyLoss()

else:

criterion = nn.BCEWithLogitsLoss()

lrs = []

best_loss = 10000

for epoch in range(epochs):

cur_lr = optimizer.param_groups[0]['lr']

print('\nEpoch=', (epoch 1), ' lr=', cur_lr)

net.train()

epoch_loss = 0

with tqdm(total=n_train, desc=f'Epoch {epoch 1}/{epochs}', unit='img') as pbar:

for batch in train_loader:

imgs = batch['image']

true_masks = batch['mask']

assert imgs.shape[1] == net.n_channels, f'Network has been defined with {net.n_channels} input channels, ' \

f'but loaded images have {imgs.shape[1]} channels. Please check that the images are loaded correctly.'

imgs = imgs.to(device=device, dtype=torch.float32)

mask_type = torch.float32 if net.n_classes == 1 else torch.long

true_masks = true_masks.to(device=device, dtype=mask_type)

masks_pred = net(imgs)

loss = criterion(masks_pred, true_masks)

epoch_loss = loss.item()

writer.add_scalar('Loss/train', loss.item(), global_step)

pbar.set_postfix(**{'loss (batch)': loss.item()})

optimizer.zero_grad()

#loss.backward()

with amp.scale_loss(loss, optimizer) as scaled_loss:

scaled_loss.backward()

optimizer.step()

pbar.update(imgs.shape[0])

global_step = 1

dataset_len = len(dataset)

a1 = dataset_len // 10 if dataset_len // 10 > 0 else 1

a2 = dataset_len / 10 if dataset_len / 10 > 0 else 1

b1 = global_step % a1

b2 = global_step % a2

if global_step % (len(dataset) // (10 * batch_size)) == 0:

val_score = eval_net(net, val_loader, device, n_val)

if net.n_classes > 1:

logging.info('Validation cross entropy: {}'.format(val_score))

writer.add_scalar('Loss/test', val_score, global_step)

else:

logging.info('Validation Dice Coeff: {}'.format(val_score))

writer.add_scalar('Dice/test', val_score, global_step)

writer.add_images('images', imgs, global_step)

if net.n_classes == 1:

writer.add_images('masks/true', true_masks, global_step)

writer.add_images('masks/pred', torch.sigmoid(masks_pred) > 0.5, global_step)

# update scheduler

scheduler.step()

lrs.append(cur_lr)

if save_cp:

try:

os.mkdir(dir_checkpoint)

logging.info('Created checkpoint directory')

except OSError:

pass

if loss.item() < best_loss:

torch.save(net.state_dict(), dir_checkpoint f'CP_epoch{epoch 1}_loss_{str(loss.item())}.pth')

best_loss = loss.item()

logging.info(f'Checkpoint {epoch 1} saved ! loss (batch) = ' str(loss.item()))

# plot lr scheduler

plt.plot(lrs , '.-', label='LambdaLR')

plt.xlabel('epoch')

plt.ylabel('LR')

plt.tight_layout()

plt.savefig('LR.png', dpi=300)

parser = argparse.ArgumentParser(description='Train the UNet on images and target masks',

formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('-g', '--gpu_id', dest='gpu_id', metavar='G', type=int, default=0, help='Number of gpu')

parser.add_argument('-u', '--unet_type', dest='unet_type', metavar='U', type=str, default='v3', help='UNet type is v1/v2/v3 (unet unet unet3 )')

parser.add_argument('-e', '--epochs', metavar='E', type=int, default=10000, help='Number of epochs', dest='epochs')

parser.add_argument('-b', '--batch-size', metavar='B', type=int, nargs='?', default=2, help='Batch size', dest='batchsize')

parser.add_argument('-l', '--learning-rate', metavar='LR', type=float, nargs='?', default=0.1, help='Learning rate', dest='lr')

parser.add_argument('-f', '--load', dest='load', type=str, default=False, help='Load model from a .pth file')

parser.add_argument('-s', '--scale', dest='scale', type=float, default=0.5, help='Downscaling factor of the images')

parser.add_argument('-v', '--validation', dest='val', type=float, default=10.0, help='Percent of the data that is used as validation (0-100)')