def scsg_diagnostics()

in diagnostics.py [0:0]

• 61 lines of code
• 11 McCabe index (conditional complexity)

def scsg_diagnostics(epoch, args, train_loader, optimizer, model, criterion):
    if (epoch == 1 or (epoch % 10) == 0) and optimizer.epoch >= optimizer.vr_from_epoch:
        nbatches = len(train_loader)

        mega_batch_size = optimizer.megabatch_size
        recalibration_interval = optimizer.recalibration_interval

        #print("interval, interval = {}".format(interval))
        optimizer.logging_pass_start()

        # Compute the snapshot
        data_buffer = []
        inner_iters = optimizer.recalibration_interval
        megabatch_size = optimizer.megabatch_size
        optimizer.recalibration_i = 0

        for batch_idx, (data, target) in enumerate(train_loader):
            batch_id = batch_idx

            if args.cuda:
                data, target = data.cuda(), target.cuda()
            data, target = Variable(data), Variable(target)

            data_buffer.append((data, target))

            # Store megabatch gradients
            def outer_closure():
                optimizer.zero_grad()
                output = model(data)
                loss = criterion(output, target)
                loss.backward()
                return loss

            loss = optimizer.step_outer_part(closure=outer_closure)

            if len(data_buffer) == megabatch_size:
                logging.info("Snapshot complete")

                for interval in range(recalibration_interval):
                    logging.info("Interval: {}, recal_i: {}".format(interval, optimizer.recalibration_i))

                    optimizer.full_grad_init()

                    # Do a full gradient calculation:
                    for inner_batch_idx, (data, target) in enumerate(train_loader):
                        if args.cuda:
                            data, target = data.cuda(), target.cuda()
                        data, target = Variable(data), Variable(target)

                        def eval_closure():
                            optimizer.zero_grad()
                            output = model(data)
                            loss = criterion(output, target)
                            loss.backward()
                            return loss

                        optimizer.full_grad_calc(closure=eval_closure)
                    logging.info("Full grad calculation finished")

                    for inner_i in range(inner_iters):
                        data, target = data_buffer[inner_i]

                        def eval_closure():
                            optimizer.zero_grad()
                            output = model(data)
                            loss = criterion(output, target)
                            loss.backward()
                            return loss

                        optimizer.logging_pass(interval, closure=eval_closure)
                    logging.info("Logging pass finished")

                    # Take a single step at the end to progress in the interval
                    data, target = data_buffer[interval]

                    def eval_closure():
                        optimizer.zero_grad()
                        output = model(data)
                        loss = criterion(output, target)
                        loss.backward()
                        return loss

                    optimizer.step_inner_part(closure=eval_closure)

                data_buffer = []
                optimizer.recalibration_i = 0
                return