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()