Update train.py

[MaddyRizvi]

Optimize GPU Utilization in Training Script using DataParallel

Description

This PR enhances GPU utilization in the training script by implementing torch.nn.DataParallel. The following improvements have been made:

• ✅ Multi-GPU Support → Enables training on multiple GPUs using torch.nn.DataParallel.
• ✅ Explicit Model-to-GPU Transfer → Moves the model to the appropriate CUDA device.
• ✅ Better Logging → Logs the number of GPUs being used for debugging and monitoring.

These changes ensure efficient parallel training without altering the existing logic of the script.

---

Changes Made

Before:

os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
logger.info(f"Work on GPU: {os.environ['CUDA_VISIBLE_DEVICES']}")

trainer = Trainer(args, load_dir=cfg.MODEL.weight_path)

After (Optimized):

os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

logger.info(f"Work on GPU: {os.environ['CUDA_VISIBLE_DEVICES']}")

trainer = Trainer(args, load_dir=cfg.MODEL.weight_path)
trainer.model.to(device)

if torch.cuda.device_count() > 1:
    trainer.model = DataParallel(trainer.model)
    logger.info(f"Using {torch.cuda.device_count()} GPUs for training")

Why This Change?:

• ✅🚀 Better Multi-GPU Utilization → Ensures all available GPUs are used for faster training.
• ✅💡 Explicit CUDA Device Assignment → Moves the model properly to the GPU.
• ✅📊 Logging Improvements → Helps track GPU usage during training.

[xh20]

it is safer move these optimized lines to the base.py class BaseTrainer, because the original BaseTrainer has used a DataParallel for the model, here I used the device (int list) instead of gpu (string)

        self.output_device = args.device[0] if type(args.device) is list else args.device

        if isintance(args.device, list) and len(args.device) > 1:
            self.model = nn.DataParallel(
                self.model,
                device_ids=args.device,
                output_device=self.output_device
            )