如何将argparse参数传递给不同方法？调用split_dataset时程序卡顿

Hey there! Let's break down your two issues with argparse and that stuck split_dataset method step by step.

First: Organizing Argparse for Different Functions #

The cleanest way to map command-line arguments to specific functions is to use argparse subcommands. This lets you group parameters by the action you want to perform (like splitting data, training a model, etc.), which keeps your code organized and avoids messy conditional checks.

Here's a concrete, reusable example:

import argparse

def split_dataset(args):
    # Use args.input_dir, args.output_train, args.output_val here
    print(f"Splitting dataset from {args.input_dir} into {args.output_train} and {args.output_val}")
    # Your actual split logic goes here

def train_model(args):
    # Use args.model_path, args.epochs, args.lr here
    print(f"Training model from {args.model_path} for {args.epochs} epochs with lr={args.lr}")
    # Your training logic goes here

def main():
    parser = argparse.ArgumentParser(description="Your ML Pipeline Tool")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Subparser for splitting dataset
    split_parser = subparsers.add_parser("split", help="Split dataset into train/val/test")
    split_parser.add_argument("--input-dir", required=True, help="Path to raw dataset")
    split_parser.add_argument("--output-train", required=True, help="Path to save training set")
    split_parser.add_argument("--output-val", required=True, help="Path to save validation set")
    split_parser.set_defaults(func=split_dataset)  # Map this subcommand to the split function

    # Subparser for training model
    train_parser = subparsers.add_parser("train", help="Train a model on the dataset")
    train_parser.add_argument("--model-path", required=True, help="Path to model checkpoint")
    train_parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
    train_parser.add_argument("--lr", type=float, default=0.001, help="Learning rate")
    train_parser.set_defaults(func=train_model)  # Map this subcommand to the train function

    args = parser.parse_args()
    args.func(args)  # Call the mapped function with parsed arguments

if __name__ == "__main__":
    main()

This structure keeps each function's parameters isolated, and you don't have to write a bunch of messy if args.split: checks. When you run the script, you'd use commands like:

python your_script.py split --input-dir ./data --output-train ./train --output-val ./val
python your_script.py train --model-path ./model.pth --epochs 20 --lr 0.0005

Second: Fixing the split_dataset Hang #

If your split_dataset method is freezing up, here are the most common culprits and actionable fixes:

• Large dataset + insufficient memory: If you're loading the entire dataset into RAM at once, it can grind your system to a halt. Instead, process the data in small batches or use a generator to load files one at a time. For example, iterate over files in the input directory, split them incrementally, and write to output without holding everything in memory.

• Slow I/O operations: Writing/reading files one by one can be glacial, especially on HDDs. Try:

  • Switching to faster storage (like an SSD)
  • Batching file writes (collect a group of files before writing to disk)
  • Using multiprocessing to handle file copying/moving in parallel (check out the multiprocessing module or concurrent.futures)

• Redundant computations or infinite loops: Double-check your split logic—are you accidentally looping forever? For example, a while loop without a proper exit condition, or recalculating split indices thousands of times unnecessarily. Add print statements or debug logs at key steps (like "Processing file X", "Writing batch to output") to track where the code gets stuck.

• Unoptimized split logic: If you're doing heavy preprocessing (like resizing images, parsing large CSV files) during the split, move that to a separate step. Keep the split function focused solely on partitioning the data, not transforming it.

• Locking issues: If multiple processes are trying to write to the same output directory at once, you might hit a file lock. Make sure each output file has a unique name, and avoid overwriting files without explicit checks.

Start with adding simple debug prints at key points in split_dataset—this will help you pinpoint exactly where it's hanging, making it way easier to fix.

内容的提问来源于stack exchange，提问作者Kim Nicoli