Function compute_labels

main.py:240–275

Pre-generate pseudo labels via network forward or uniformly assignment

(dataloader, model, class_num, data_len, init_via_forward=False)

    return losses.avg

def compute_labels(dataloader, model, class_num, data_len, init_via_forward=False):
    '''Pre-generate pseudo labels via network forward or uniformly assignment'''
    if args.verbose:
        logger.info('Compute labels')
    batch_time = AverageMeter()
    data_time = AverageMeter()
    end = time.time()

    if init_via_forward:
        model.eval()
        label_list = []
        for i, inputs in enumerate(dataloader):
            data_time.update(time.time() - end)

            inputs = inputs.cuda()
            with torch.no_grad():
                output = model(inputs)
            batch_time.update(time.time() - end)

            label = output.argmax(dim=1).cpu()
            if i == 0:
                label_list = label
            else:
                label_list = torch.cat([label_list, label], dim=0)
            if args.verbose and (i % 100) == 0:
                logger.info('{0}/{1}\t'
                      'Time: {batch_time.val:.3f} ({batch_time.avg:.3f})'
                      'Data: {data_time.val:.3f} ({data_time.avg:.3f})\t'
                      .format(i, len(dataloader), batch_time=batch_time, data_time=data_time))
            end = time.time()
        label_list = label_list.numpy()
        model.train()
    else:
        label_list = np.array([int(np.random.uniform(0, class_num)) for _ in range(data_len)])

    return label_list

if __name__ == '__main__':
    main()