<commit_before>"""
Modified from: https://github.com/pytorch/examples/blob/master/imagenet/main.py
"""
import shutil
import time

import torch
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed

try:
    from torch.utils.data._utils.collate import default_collate
except ModuleNotFoundError:
    # import from older versions of pytorch
    from torch.utils.data.dataloader import default_collate


def epoch_training(train_loader, model, criterion, optimizer, epoch, n_gpus=None, print_frequency=1, regularized=False,
                   print_gpu_memory=False, vae=False, scaler=None):
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, losses],
        prefix="Epoch: [{}]".format(epoch))

    use_amp = scaler is not None

    # switch to train mode
    model.train()

    end = time.time()
    for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        if n_gpus:
            # torch.cuda.empty_cache()
            if print_gpu_memory:
                for i_gpu in range(n_gpus):
                    print("Memory allocated (device {}):".format(i_gpu),
                          human_readable_size(torch.cuda.memory_allocated(i_gpu)))
                    print("Max memory allocated (device {}):".format(i_gpu),
                          human_readable_size(torch.cuda.max_memory_allocated(i_gpu)))
                    print("Memory cached (device {}):".format(i_gpu),
                          human_readable_size(torch.cuda.memory_cached(i_gpu)))
                    print("Max memory cached (device {}):".format(i_gpu),
                          human_readable_size(torch.cuda.max_memory_cached(i_gpu)))

        optimizer.zero_grad()
        loss, batch_size = batch_loss(model, images, target, criterion, n_gpus=n_gpus, regularized=regularized,
                                      vae=vae, use_amp=use_amp)
        if n_gpus:
            torch.cuda.empty_cache()

        # measure accuracy and record loss
        losses.update(loss.item(), batch_size)

        if scaler:
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            # compute gradient and do step
            loss.backward()
            optimizer.step()

        del loss

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % print_frequency == 0:
            progress.display(i)
    return losses.avg


def batch_loss(model, images, target, criterion, n_gpus=0, regularized=False, vae=False, use_amp=None):
    if n_gpus is not None:
        images = images.cuda()
        target = target.cuda()
    # compute output
    if use_amp:
        from torch.cuda.amp import autocast
        with autocast():
            return _batch_loss(model, images, target, criterion, regularized=regularized, vae=vae)
    else:
        return _batch_loss(model, images, target, criterion, regularized=regularized, vae=vae)


def _batch_loss(model, images, target, criterion, regularized=False, vae=False):
    output = model(images)
    batch_size = images.size(0)
    if regularized:
        try:
            output, output_vae, mu, logvar = output<commit_msg>ENH: debug<commit_after>"""
Modified from: https://github.com/pytorch/examples/blob/master/imagenet/main.py
"""
import shutil
import time

import torch
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed

try:
    from torch.utils.data._utils.collate import default_collate
except ModuleNotFoundError:
    # import from older versions of pytorch
    from torch.utils.data.dataloader import default_collate


def epoch_training(train_loader, model, criterion, optimizer, epoch, n_gpus=None, print_frequency=1, regularized=False,
                   print_gpu_memory=False, vae=False, scaler=None):
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, losses],
        prefix="Epoch: [{}]".format(epoch))

    use_amp = scaler is not None

    # switch to train mode
    model.train()

    end = time.time()
    for i, (images, target) in enumerate(train_loader):
        print("epoch_training:", target.size())
        # measure data loading time
        data_time.update(time.time() - end)

        if n_gpus:
            # torch.cuda.empty_cache()
            if print_gpu_memory:
                for i_gpu in range(n_gpus):
                    print("Memory allocated (device {}):".format(i_gpu),
                          human_readable_size(torch.cuda.memory_allocated(i_gpu)))
                    print("Max memory allocated (device {}):".format(i_gpu),
                          human_readable_size(torch.cuda.max_memory_allocated(i_gpu)))
                    print("Memory cached (device {}):".format(i_gpu),
                          human_readable_size(torch.cuda.memory_cached(i_gpu)))
                    print("Max memory cached (device {}):".format(i_gpu),
                          human_readable_size(torch.cuda.max_memory_cached(i_gpu)))

        optimizer.zero_grad()
        loss, batch_size = batch_loss(model, images, target, criterion, n_gpus=n_gpus, regularized=regularized,
                                      vae=vae, use_amp=use_amp)
        if n_gpus:
            torch.cuda.empty_cache()

        # measure accuracy and record loss
        losses.update(loss.item(), batch_size)

        if scaler:
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            # compute gradient and do step
            loss.backward()
            optimizer.step()

        del loss

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % print_frequency == 0:
            progress.display(i)
    return losses.avg


def batch_loss(model, images, target, criterion, n_gpus=0, regularized=False, vae=False, use_amp=None):
    print("batch_loss:", target.size())
    if n_gpus is not None:
        images = images.cuda()
        target = target.cuda()
    # compute output
    if use_amp:
        from torch.cuda.amp import autocast
        with autocast():
            return _batch_loss(model, images, target, criterion, regularized=regularized, vae=vae)
    else:
        return _batch_loss(model, images, target, criterion, regularized=regularized, vae=vae)


def _batch_loss(model, images, target, criterion, regularized=False, vae=False):
    print("_batch_loss:", target.size())
    output = model(images)
    batch_size = images.size(0)
    if regularized:
        try:
            output, output_vae, mu, logvar = output
{'fraction_non_alphanumeric': 0.1070518266779949, 'fraction_numerical': 0.002548853016142736, 'mean_word_length': 2.4269771955361477, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1674425', 'n_tokens_mistral': 2165, 'n_tokens_neox': 2079, 'n_words': 468}