DDP question to trainer fit test predict and savings

[felixdittrich92]

Hi,
i have a short question to DDP is it the correct way to combine all like fit / test / predict and some savings in the main script or is there a need to split into parts ? It works but feels a bit coriuos
my example:

if __name__ == '__main__':

    warnings.filterwarnings("ignore")

    pl.seed_everything(42, workers=True)

    print('annotations file and image folder have to be in the same parent folder')

    parser = argparse.ArgumentParser(description='Document Layout Detection Trainer')
    parser.add_argument("--path", help='path to generated images', type=str, required=False, default='/PublayNetDataset') #set to true
    parser.add_argument("--train_ann_file_name", help='name of annotations file', type=str, required=False, default='train.json')
    parser.add_argument("--train_image_folder_name", help='name of image folder', type=str, required=False, default='train')
    parser.add_argument("--val_ann_file_name", help='name of annotations file', type=str, required=False, default='val.json')
    parser.add_argument("--val_image_folder_name", help='name of image folder', type=str, required=False, default='val')
    parser.add_argument("--test_image_folder_name", help='name of image folder', type=str, required=False, default='test')
    parser.add_argument("--epochs", help='how many epochs to train the model',type=int, required=False, default=50)
    parser.add_argument("--batch_size", help='how big are a batch',type=int, required=False, default=8)
    parser.add_argument("--data_limit", help='how many data to use',type=int, required=False, default=0)
    parser.add_argument("--worker", help='how many threads for the Dataloader',type=int, required=False, default=0)
    parser.add_argument("--learning_rate", help='the learning rate for the optimizer',type=float, required=False, default=0.0001)
    parser.add_argument("--gradient_clip", help='float for gradient clipping',type=float, required=False, default=0.5)
    parser.add_argument("--visualize_random_example", help='if true show an example from train set',type=__check_for_boolean_value, required=False, default=False)

    args = parser.parse_args()
    path = args.path
    train_ann_file_name = args.train_ann_file_name
    train_image_folder_name = args.train_image_folder_name
    val_ann_file_name = args.val_ann_file_name
    val_image_folder_name = args.val_image_folder_name
    test_image_folder_name = args.test_image_folder_name
    epochs = args.epochs
    batch_size = args.batch_size
    data_limit = args.data_limit
    worker = args.worker
    learning_rate = args.learning_rate
    gradient_clip = args.gradient_clip
    visualize_random_example = args.visualize_random_example

    # resource handling
    if torch.cuda.device_count() >= 1:
        batch_size = int(batch_size / torch.cuda.device_count())
        accelerator = 'ddp'
        sync = True
    else:
        accelerator = None
        sync = False

    ### Data Part

    if data_limit == 0:
        data_limit = None

    os.makedirs('layout_detection_model_files', exist_ok=True)

    feature_extractor = DetrFeatureExtractor.from_pretrained("facebook/detr-resnet-50", do_resize=True, size=600, max_size=1333)
    feature_extractor.save_pretrained('layout_detection_model_files/transformer_model/')
    print('feature extractor saved succesful')

    data_module = COCODatasetLoader(path,
                                    train_ann_file_name,
                                    train_image_folder_name,
                                    val_ann_file_name,
                                    val_image_folder_name,
                                    test_image_folder_name,
                                    feature_extractor,
                                    batch_size,
                                    worker,
                                    train_collator=collate_fn_train,
                                    test_collator=collate_fn_test,
                                    data_limit=data_limit)
    data_module.setup()
    if visualize_random_example:
        index = np.random.choice(len(data_module.train_dataset))
        data_module.visualize_example(index)

    train = data_module.train_dataloader()
    val = data_module.val_dataloader()
    test = data_module.test_dataloader()
    coco_val_dataset = data_module.get_val_coco_text_dataset()
    coco_evaluator = CocoEvaluator(coco_val_dataset, ['bbox'])
    print('Coco Evaluator created')

    cats = coco_val_dataset.cats
    id2label = {k: v['name'] for k,v in cats.items()}
    # save id2label file
    with open('layout_detection_model_files/id2label.json', 'w') as fp:
        json.dump(id2label, fp)
    print('Labels file for Onnx saved successful')
    print(f'IDs : {id2label}')

    ### Model Part
    layout_detection_model = LayoutDetectionModel(lr=learning_rate, id2label=id2label, feature_extractor=feature_extractor, coco_evaluator=coco_evaluator, sync=sync)

    ### Callback Part

    checkpoint_callback = ModelCheckpoint(
        dirpath="layout_detection_model_files/checkpoints",
        filename="best-checkpoint",
        save_top_k=1,
        verbose=True,
        monitor="val_loss",
        mode="min"
    )

    logger = TensorBoardLogger(save_dir="layout_detection_model_files/Lightning_logs", name="Layout_Detection")

    early_stopping_callback = EarlyStopping(
        monitor="val_loss",
        min_delta=0.001,
        patience=5,
        check_finite=True,
        verbose=True
    )

    lr_monitor = LearningRateMonitor(logging_interval='epoch')

    ### Training Part

    trainer = pl.Trainer(logger=logger,
                         weights_summary="full",
                         # only if gpu mem is overheaded -> needs much more train time
                         move_metrics_to_cpu=False,
                         benchmark=True,
                         gradient_clip_val=gradient_clip,  # set to 0.5 to avoid exploding gradients
                         stochastic_weight_avg=True,
                         callbacks=[
                             checkpoint_callback,
                             early_stopping_callback,
                             lr_monitor
                             ],
                         max_epochs=epochs,
                         gpus=torch.cuda.device_count(),
                         accelerator=accelerator,
                         precision=32, # do not change for this model training
                         accumulate_grad_batches=1,  # optimizer step after every n batches  -> better gpu mem usage / model specific
                         progress_bar_refresh_rate=20,
                        # profiler='pytorch', # only for debug
                         )

    trainer.fit(layout_detection_model, train, val)
    trainer.test(layout_detection_model, test)

    time.sleep(2) # short delay
    trained_model = layout_detection_model.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)
    trained_model.eval()
    trained_model.freeze()

    ### Saving Part

    # ----------------------------------
    # PyTorch Model - full
    # ----------------------------------
    try:
        torch.save(trained_model, "layout_detection_model_files/torch_layout_detection_model.pt")
        print('Torch model saved successful')
    except Exception as e:
        print('Cannot export as PyTorch Format -- Error : ' + str(e))

    # ----------------------------------
    # PyTorch Model - state dict
    # ----------------------------------
    try:
        torch.save(trained_model.state_dict(), "layout_detection_model_files/torch_layout_detection_model_state_dict.pt")
        print('Torch model state dict saved successful')
    except Exception as e:
        print('Cannot export as PyTorch Format with state dict -- Error : ' + str(e))

    # ----------------------------------
    # onnx
    # ----------------------------------
    try:
        input_batch = next(iter(val))
        input_sample = {
            "pixel_values": input_batch["pixel_values"][0].unsqueeze(0),
            "pixel_mask": input_batch["pixel_mask"][0].unsqueeze(0)
        }
        values = input_sample['pixel_values']
        mask = input_sample['pixel_mask']
        file_path = "layout_detection_model_files/torch_layout_detection_model.onnx"
        torch.onnx.export(trained_model, (values, mask), file_path,
                              input_names=['pixel_values', 'pixel_mask'],
                              output_names=['logits'],
                              dynamic_axes={'pixel_values': {0: 'batch_size', 1: 'channels', 2: 'width', 3: 'height'}, 'pixel_mask': {0: 'batch_size', 1: 'width', 2: 'height'}, 'logits': {0: 'batch_size'}},
                              export_params=True, opset_version=11,
                              enable_onnx_checker=True, verbose=False)
        print('Onnx model saved successful')
        print('Start model quantization')
        model_quant = "layout_detection_model_files/torch_layout_detection_model.quant.onnx"
        quantized_model = quantize_qat(file_path, model_quant)
        print('Quantization succesfull')
    except Exception as e:
        print('Cannot export as ONNX Format -- Error : ' + str(e))


    # Predictions
    model = layout_detection_model.load_from_checkpoint(checkpoint_path=trainer.checkpoint_callback.best_model_path)
    preds = trainer.predict(model, test, return_predictions=True)
    print(preds)

Would be great to get a short feedback or exist a full example for some scripts without the LightningCLI ?
Best regards
Felix