Saltar al contenido

Referencia para ultralytics/utils/autobatch.py

Nota

Este archivo est谩 disponible en https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/utils/autobatch .py. Si detectas alg煤n problema, por favor, ayuda a solucionarlo contribuyendo con una Pull Request 馃洜锔. 隆Gracias 馃檹!



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

Comprueba YOLO el tama帽o del lote de entrenamiento utilizando la funci贸n autobatch().

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
model Module

YOLO modelo para comprobar el tama帽o del lote.

necesario
imgsz int

Tama帽o de la imagen utilizada para el entrenamiento.

640
amp bool

Si es Verdadero, utiliza la precisi贸n mixta autom谩tica (PAM) para el entrenamiento.

True

Devuelve:

Tipo Descripci贸n
int

Tama帽o 贸ptimo del lote calculado con la funci贸n autobatch().

C贸digo fuente en 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)

Estima autom谩ticamente el mejor tama帽o de lote YOLO para utilizar una fracci贸n de la memoria CUDA disponible.

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
model module

YOLO modelo para calcular el tama帽o del lote.

necesario
imgsz int

El tama帽o de imagen utilizado como entrada para el modelo YOLO . Por defecto es 640.

640
fraction float

La fracci贸n de memoria CUDA disponible a utilizar. Por defecto es 0,60.

0.6
batch_size int

El tama帽o de lote por defecto que se utilizar谩 si se detecta un error. Por defecto es 16.

batch

Devuelve:

Tipo Descripci贸n
int

El tama帽o 贸ptimo del lote.

C贸digo fuente en 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





Creado 2023-11-12, Actualizado 2024-05-18
Autores: glenn-jocher (4), Burhan-Q (1), Laughing-q (1)