Zum Inhalt springen

Referenz fĂŒr ultralytics/utils/autobatch.py

Hinweis

Diese Datei ist verfĂŒgbar unter https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/utils/autobatch .py. Wenn du ein Problem entdeckst, hilf bitte mit, es zu beheben, indem du einen Pull Request đŸ› ïž einreichst. Vielen Dank 🙏!



ultralytics.utils.autobatch.check_train_batch_size(model, imgsz=640, amp=True)

ÜberprĂŒfe die YOLO TrainingsstapelgrĂ¶ĂŸe mit der Funktion autobatch().

Parameter:

Name Typ Beschreibung Standard
model Module

YOLO Modell fĂŒr die ÜberprĂŒfung der LosgrĂ¶ĂŸe.

erforderlich
imgsz int

BildgrĂ¶ĂŸe, die fĂŒr das Training verwendet wird.

640
amp bool

Wenn True, verwende automatische gemischte PrĂ€zision (AMP) fĂŒr die Ausbildung.

True

Retouren:

Typ Beschreibung
int

Optimale BatchgrĂ¶ĂŸe, die mit der Funktion autobatch() berechnet wird.

Quellcode in ultralytics/utils/autobatch.py
def check_train_batch_size(model, imgsz=640, amp=True):
    """
    Check YOLO training batch size using the autobatch() function.

    Args:
        model (torch.nn.Module): YOLO model to check batch size for.
        imgsz (int): Image size used for training.
        amp (bool): If True, use automatic mixed precision (AMP) for training.

    Returns:
        (int): Optimal batch size computed using the autobatch() function.
    """

    with torch.cuda.amp.autocast(amp):
        return autobatch(deepcopy(model).train(), imgsz)  # compute optimal batch size



ultralytics.utils.autobatch.autobatch(model, imgsz=640, fraction=0.6, batch_size=DEFAULT_CFG.batch)

SchĂ€tze automatisch die beste YOLO StapelgrĂ¶ĂŸe, um einen Bruchteil des verfĂŒgbaren CUDA Speichers zu nutzen.

Parameter:

Name Typ Beschreibung Standard
model module

YOLO Modell, fĂŒr das die LosgrĂ¶ĂŸe berechnet wird.

erforderlich
imgsz int

Die BildgrĂ¶ĂŸe, die als Eingabe fĂŒr das Modell YOLO verwendet wird. Der Standardwert ist 640.

640
fraction float

Der Anteil des verfĂŒgbaren CUDA Speichers, der verwendet werden soll. Der Standardwert ist 0,60.

0.6
batch_size int

Die Standard-StapelgrĂ¶ĂŸe, die verwendet wird, wenn ein Fehler festgestellt wird. Der Standardwert ist 16.

batch

Retouren:

Typ Beschreibung
int

Die optimale LosgrĂ¶ĂŸe.

Quellcode in ultralytics/utils/autobatch.py
def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch):
    """
    Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.

    Args:
        model (torch.nn.module): YOLO model to compute batch size for.
        imgsz (int, optional): The image size used as input for the YOLO model. Defaults to 640.
        fraction (float, optional): The fraction of available CUDA memory to use. Defaults to 0.60.
        batch_size (int, optional): The default batch size to use if an error is detected. Defaults to 16.

    Returns:
        (int): The optimal batch size.
    """

    # Check device
    prefix = colorstr("AutoBatch: ")
    LOGGER.info(f"{prefix}Computing optimal batch size for imgsz={imgsz}")
    device = next(model.parameters()).device  # get model device
    if device.type == "cpu":
        LOGGER.info(f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}")
        return batch_size
    if torch.backends.cudnn.benchmark:
        LOGGER.info(f"{prefix} ⚠ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}")
        return batch_size

    # Inspect CUDA memory
    gb = 1 << 30  # bytes to GiB (1024 ** 3)
    d = str(device).upper()  # 'CUDA:0'
    properties = torch.cuda.get_device_properties(device)  # device properties
    t = properties.total_memory / gb  # GiB total
    r = torch.cuda.memory_reserved(device) / gb  # GiB reserved
    a = torch.cuda.memory_allocated(device) / gb  # GiB allocated
    f = t - (r + a)  # GiB free
    LOGGER.info(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free")

    # Profile batch sizes
    batch_sizes = [1, 2, 4, 8, 16]
    try:
        img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]
        results = profile(img, model, n=3, device=device)

        # Fit a solution
        y = [x[2] for x in results if x]  # memory [2]
        p = np.polyfit(batch_sizes[: len(y)], y, deg=1)  # first degree polynomial fit
        b = int((f * fraction - p[1]) / p[0])  # y intercept (optimal batch size)
        if None in results:  # some sizes failed
            i = results.index(None)  # first fail index
            if b >= batch_sizes[i]:  # y intercept above failure point
                b = batch_sizes[max(i - 1, 0)]  # select prior safe point
        if b < 1 or b > 1024:  # b outside of safe range
            b = batch_size
            LOGGER.info(f"{prefix}WARNING ⚠ CUDA anomaly detected, using default batch-size {batch_size}.")

        fraction = (np.polyval(p, b) + r + a) / t  # actual fraction predicted
        LOGGER.info(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅")
        return b
    except Exception as e:
        LOGGER.warning(f"{prefix}WARNING ⚠ error detected: {e},  using default batch-size {batch_size}.")
        return batch_size





Erstellt am 2023-11-12, Aktualisiert am 2023-11-25
Autoren: glenn-jocher (3), Laughing-q (1)