if command == '--Help':
print("Check the documentation.")
elif command == "--Train" or command == "--Predict":
try:
architecture = sys.argv[2]
except IndexError:
print(
"An Index error has occured! Check the documentation to make sure your passing all the required arguments."
)
if command == "--Train":
if architecture == '--UNet':
UNet = UNet()
UNet.train()
elif architecture == '--UNet++':
UNetPP = UNetPP()
UNetPP.train()
else:
raise Exception(
"You have passed an invalid argument.\nCheck the documentation for the allowed arguments."
)
elif command == "--Predict":
if architecture == '--UNet':
UNet = UNet()
UNet.predict()
elif architecture == '--UNet++':
UNetPP = UNetPP()
UNetPP.predict()
num_batches = len(dataloader.train_objs)
num_batches_val = len(dataloader.val_objs)
num_lr_increases = num_batches * num_epochs
num_lr_increase = 0
if torch.cuda.is_available():
model.cuda(params['device'])
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=momentum)
start = time.time()
for num_epoch in range(num_epochs):
dataloader.on_epoch_start(params)
for num_batch in range(num_batches):
model.train()
X, Y = dataloader.get_mini_batch(num_batch,
params,
mode='train',
weight=False,
data_augmentation=True)
lr = min_lr * 3.**((float(num_lr_increase) / num_lr_increases) *
(np.log(max_lr / min_lr) / np.log(3.)))
set_lr(optimizer, lr)
optimizer.zero_grad()
out = model(X)
Y = crop_center_Y(out, Y)
loss = 1 - DICE_coefficient(out, Y, zero_division_safe=True)
loss.backward()
optimizer.step()