Exemplo n.º 1
0
def set_model(args, n_data):
    # set the model
    if args.model == 'c3d':
        model = C3D(with_classifier=False)
    elif args.model == 'r3d':
        model = R3DNet(layer_sizes=(1, 1, 1, 1), with_classifier=False)
    elif args.model == 'r21d':
        model = R2Plus1DNet(layer_sizes=(1, 1, 1, 1), with_classifier=False)

    if args.intra_neg:
        contrast = NCEAverage(args.feat_dim, n_data, args.nce_k, args.nce_t,
                              args.nce_m, args.softmax)
    else:
        contrast = NCEAverage_ori(args.feat_dim, n_data, args.nce_k,
                                  args.nce_t, args.nce_m, args.softmax)

    criterion_1 = NCESoftmaxLoss() if args.softmax else NCECriterion(n_data)
    criterion_2 = NCESoftmaxLoss() if args.softmax else NCECriterion(n_data)

    # GPU mode
    model = model.cuda()
    contrast = contrast.cuda()
    criterion_1 = criterion_1.cuda()
    criterion_2 = criterion_2.cuda()
    cudnn.benchmark = True

    return model, contrast, criterion_1, criterion_2
Exemplo n.º 2
0
                                device_ids=range(torch.cuda.device_count()))
    cudnn.benchmark = True

# Model
if args.test_only or len(args.resume) > 0:
    # Load checkpoint.
    print('==> Resuming from checkpoint..')
    assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
    checkpoint = torch.load('./checkpoint/' + args.resume)
    net.load_state_dict(checkpoint['net'])
    lemniscate = checkpoint['lemniscate']
    best_acc = checkpoint['acc']
    start_epoch = checkpoint['epoch']

# define loss function
criterion = NCECriterion()

net.to(device)
lemniscate.to(device)
criterion.to(device)

if args.test_only:
    acc = kNN(0, net, lemniscate, trainloader, testloader, 200, args.nce_t, 1)
    sys.exit(0)

optimizer = optim.SGD(net.parameters(),
                      lr=args.lr,
                      momentum=0.9,
                      weight_decay=5e-4)

Exemplo n.º 3
0
def main():
    global args, best_prec1
    args = parser.parse_args()

    args.distributed = args.world_size > 1

    if args.distributed:
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                world_size=args.world_size)

    # create model
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=> creating model '{}'".format(args.arch))
        model = models.__dict__[args.arch](low_dim=args.low_dim)

    if not args.distributed:
        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
            model.features = torch.nn.DataParallel(model.features)
            model.cuda()
        else:
            model = torch.nn.DataParallel(model).cuda()
    else:
        model.cuda()
        model = torch.nn.parallel.DistributedDataParallel(model)


    # Data loading code
    traindir = os.path.join(args.data, 'train')
    valdir = os.path.join(args.data, 'val')
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    train_dataset = datasets.ImageFolderInstance(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(224, scale=(0.2,1.)),
            transforms.RandomGrayscale(p=0.2),
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]))

    if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    else:
        train_sampler = None

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=train_sampler)

    val_loader = torch.utils.data.DataLoader(
        datasets.ImageFolderInstance(valdir, transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            normalize,
        ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    # define lemniscate and loss function (criterion)
    ndata = train_dataset.__len__()
    if args.nce_k > 0:
        lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t, args.nce_m).cuda()
        criterion = NCECriterion(ndata).cuda()
    else:
        lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t, args.nce_m).cuda()
        criterion = nn.CrossEntropyLoss().cuda()

    optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            model.load_state_dict(checkpoint['state_dict'])
            lemniscate = checkpoint['lemniscate']
            optimizer.load_state_dict(checkpoint['optimizer'])
            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    cudnn.benchmark = True

    if args.evaluate:
        kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
        return

    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            train_sampler.set_epoch(epoch)
        adjust_learning_rate(optimizer, epoch)

        # train for one epoch
        train(train_loader, model, lemniscate, criterion, optimizer, epoch)

        # evaluate on validation set
        prec1 = NN(epoch, model, lemniscate, train_loader, val_loader)

        # remember best prec@1 and save checkpoint
        is_best = prec1 > best_prec1
        best_prec1 = max(prec1, best_prec1)
        save_checkpoint({
            'epoch': epoch + 1,
            'arch': args.arch,
            'state_dict': model.state_dict(),
            'lemniscate': lemniscate,
            'best_prec1': best_prec1,
            'optimizer' : optimizer.state_dict(),
        }, is_best)
    # evaluate KNN after last epoch
    kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
Exemplo n.º 4
0
def main():
    global args, best_prec1
    args = parser.parse_args()

    args.distributed = args.world_size > 1

    if args.distributed:
        dist.init_process_group(backend=args.dist_backend,
                                init_method=args.dist_url,
                                world_size=args.world_size)

    # create model
    # if args.pretrained:
    #     print("=> using pre-trained model '{}'".format(args.arch))
    #     model = models.__dict__[args.arch](pretrained=True, finetune=args.finetune, low_dim= args.low_dim)
    # else:
    #     print("=> creating model '{}'".format(args.arch))
    #
    #     model = models.__dict__[args.arch](low_dim=args.low_dim)

    # Data loading code

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    # train_dataset = datasets.CombinedMaskDataset(
    #     other_data_path = '/home/saschaho/Simcenter/found_label_imgs',
    #     csv_root_folder='/home/saschaho/Simcenter/Floor_Elevation_Data/Streetview_Irma/Streetview_Irma/images',
    #     data_csv='/home/saschaho/Simcenter/Building_Information_Prediction/all_bims_train.csv',
    #     transform = transforms.Compose([
    #         transforms.RandomResizedCrop(224, scale=(0.2,1.)),
    #         transforms.RandomGrayscale(p=0.2),
    #         transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
    #         transforms.RandomHorizontalFlip(),
    #         transforms.ToTensor(),
    #         normalize,
    #     ]),attribute = 'first_floor_elevation_ft', mask_images=True)

    # val_dataset = datasets.CombinedMaskDataset(
    #         csv_root_folder='/home/saschaho/Simcenter/Floor_Elevation_Data/Streetview_Irma/Streetview_Irma/images',
    #         data_csv='/home/saschaho/Simcenter/Building_Information_Prediction/all_bims_val.csv',
    #     transform=transforms.Compose([
    #         transforms.Resize(256),
    #         transforms.CenterCrop(224),
    #         transforms.ToTensor(),
    #         normalize,
    #     ]),
    #attribute = 'first_floor_elevation_ft', mask_images=True)

    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224, scale=(0.3, 1.)),
        transforms.RandomGrayscale(p=0.5),
        transforms.ColorJitter(0.5, 0.5, 0.5, 0.5),
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.ToTensor(), normalize
    ])

    val_transform = transforms.Compose(
        [transforms.Resize((224, 224)),
         transforms.ToTensor(), normalize])

    train_dataset = First_Floor_Binary(args.attribute_name,
                                       args.train_data,
                                       args.image_folder,
                                       transform=train_transform,
                                       regression=args.regression,
                                       mask_buildings=args.mask_buildings,
                                       softmask=args.softmask)
    val_dataset = First_Floor_Binary(args.attribute_name,
                                     args.val_data,
                                     args.image_folder,
                                     transform=val_transform,
                                     regression=args.regression,
                                     mask_buildings=args.mask_buildings,
                                     softmask=args.softmask)

    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=args.batch_size,
                                               shuffle=True,
                                               num_workers=args.workers,
                                               pin_memory=True)

    val_loader = torch.utils.data.DataLoader(val_dataset,
                                             batch_size=args.batch_size,
                                             shuffle=False,
                                             num_workers=args.workers,
                                             pin_memory=True)

    model = ResidualAttentionModel_92_Small(args.low_dim, dropout=False)
    model = torch.nn.DataParallel(model).cuda()

    print('Train dataset instances: {}'.format(len(train_loader.dataset)))
    print('Val dataset instances: {}'.format(len(val_loader.dataset)))
    # define lemniscate and loss function (criterion)
    ndata = train_dataset.__len__()
    if args.nce_k > 0:
        lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t,
                                args.nce_m).cuda()
        criterion = NCECriterion(ndata).cuda()
    else:
        lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t,
                                   args.nce_m).cuda()
        criterion = nn.CrossEntropyLoss().cuda()

    optimizer = torch.optim.SGD(model.parameters(),
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    #optimizer = RAdam(model.parameters())

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']

            keyname = [
                keyname for keyname in model.state_dict().keys()
                if 'fc.weight' in keyname
            ][0]
            lat_vec_len_model = model.state_dict()[keyname].shape[0]
            lat_vec_len_checkpoint = checkpoint['state_dict'][keyname].shape[0]

            low_dim_differ = False
            if lat_vec_len_model != lat_vec_len_checkpoint:
                low_dim_differ = True
                print(
                    'Warning: Latent vector sizes do not match. Assuming finetuning'
                )
                print(
                    'Lemniscate will be trained from scratch with new optimizer.'
                )
                del checkpoint['state_dict'][keyname]
                del checkpoint['state_dict'][keyname.replace('weight', 'bias')]

            missing_keys, unexpected_keys = model.load_state_dict(
                checkpoint['state_dict'], strict=False)
            if len(missing_keys) or len(unexpected_keys):
                print('Warning: Missing or unexpected keys found.')
                print('Missing: {}'.format(missing_keys))
                print('Unexpected: {}'.format(unexpected_keys))

            if not low_dim_differ:
                # The memory bank will be trained from scratch if
                # the low dim is different. Maybe later repopulated
                lemniscate = checkpoint['lemniscate']
                optimizer.load_state_dict(checkpoint['optimizer'])
            print("=> loaded checkpoint '{}' (epoch {})".format(
                args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    cudnn.benchmark = True

    if args.evaluate:
        kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
        return

    for epoch in range(args.start_epoch, args.epochs):
        # if args.distributed:
        #     train_sampler.set_epoch(epoch)
        #adjust_learning_rate(optimizer, epoch)

        # train for one epoch
        train(train_loader, model, lemniscate, criterion, optimizer, epoch)

        # evaluate on validation set
        prec1 = NN(epoch, model, lemniscate, train_loader, val_loader)

        # remember best prec@1 and save checkpoint
        is_best = prec1 > best_prec1
        best_prec1 = max(prec1, best_prec1)
        save_checkpoint(
            {
                'epoch': epoch + 1,
                'state_dict': model.state_dict(),
                'lemniscate': lemniscate,
                'best_prec1': best_prec1,
                'optimizer': optimizer.state_dict(),
            }, is_best, args.name)
    # evaluate KNN after last epoch
    kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
Exemplo n.º 5
0
def main():
    global args, best_prec1, best_prec1_past, best_prec1_future
    args = parser.parse_args()

    args.distributed = args.world_size > 1

    if args.distributed:
        dist.init_process_group(backend=args.dist_backend,
                                init_method=args.dist_url,
                                world_size=args.world_size)

    # create model
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True)
    else:
        print("=> creating model '{}'".format(args.arch))
        model = models.__dict__[args.arch](low_dim=args.low_dim)

    if not args.distributed:
        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
            model.features = torch.nn.DataParallel(model.features)
            model.to(get_device(args.gpu))
        else:
            model = torch.nn.DataParallel(model).to(get_device(args.gpu))
    else:
        model.to(get_device(args.gpu))
        model = torch.nn.parallel.DistributedDataParallel(model)

    # Data loading code
    traindir = os.path.join(args.data, 'train')
    valdir = os.path.join(args.data, 'val')
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    train_dataset = Dataset(traindir, n_frames)
    val_dataset = Dataset(valdir, n_frames)

    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,  #(train_sampler is None), 
        num_workers=args.workers)

    val_loader = torch.utils.data.DataLoader(val_dataset,
                                             batch_size=args.batch_size,
                                             shuffle=False,
                                             num_workers=args.workers)

    # define lemniscate and loss function (criterion)
    ndata = train_dataset.__len__()
    if args.nce_k > 0:
        lemniscate = NCEAverage(args.gpu, args.low_dim, ndata, args.nce_k,
                                args.nce_t,
                                args.nce_m).to(get_device(args.gpu))
        criterion = NCECriterion(ndata).to(get_device(args.gpu))
    else:
        lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t,
                                   args.nce_m).to(get_device(args.gpu))
        criterion = nn.CrossEntropyLoss().to(get_device(args.gpu))

    optimizer = torch.optim.SGD(model.parameters(),
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            model.load_state_dict(checkpoint['state_dict'])
            lemniscate = checkpoint['lemniscate']
            optimizer.load_state_dict(checkpoint['optimizer'])
            print("=> loaded checkpoint '{}' (epoch {})".format(
                args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    cudnn.benchmark = True

    if args.evaluate:
        kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
        return

    for epoch in range(args.start_epoch, args.epochs):

        adjust_learning_rate(optimizer, epoch)

        # train for one epoch
        train(train_loader, model, lemniscate, criterion, optimizer, epoch)

        # evaluate on validation set
        prec1, prec1_past, prec1_future = NN(epoch, model, lemniscate,
                                             train_loader, val_loader)

        add_epoch_score('epoch_scores.txt', epoch, prec1)
        add_epoch_score('epoch_scores_past.txt', epoch, prec1_past)
        add_epoch_score('epoch_scores_future.txt', epoch, prec1_future)

        # Sascha: This is a bug because it seems prec1 or best_prec1 is a vector at some point with
        # more than one entry
        # remember best prec@1 and save checkpoint
        is_best = prec1 > best_prec1
        best_prec1 = max(prec1, best_prec1)
        save_checkpoint(
            {
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'lemniscate': lemniscate,
                'best_prec1': best_prec1,
                'optimizer': optimizer.state_dict(),
            }, is_best, epoch)

        is_best_past = prec1_past > best_prec1_past
        best_prec1_past = max(prec1_past, best_prec1_past)
        save_checkpoint(
            {
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'lemniscate': lemniscate,
                'best_prec1_past': best_prec1_past,
                'optimizer': optimizer.state_dict(),
            },
            is_best_past,
            epoch,
            best_mod='_past')

        is_best_future = prec1_future > best_prec1_future
        best_prec1_future = max(prec1_future, best_prec1_future)
        save_checkpoint(
            {
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'lemniscate': lemniscate,
                'best_prec1_future': best_prec1_future,
                'optimizer': optimizer.state_dict(),
            },
            is_best_future,
            epoch,
            best_mod='_future')
    # evaluate KNN after last epoch
    kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
scheduler = get_scheduler(optimizer, len(trainloader), args)

# Model
if args.test_only or len(args.resume) > 0:
    # Load checkpoint.
    print('==> Resuming from checkpoint..')
    checkpoint = torch.load(args.resume)
    net.load_state_dict(checkpoint['net'])
    lemniscate = checkpoint['lemniscate']
    best_acc1 = checkpoint['acc']
    start_epoch = checkpoint['epoch']

# define loss function
if hasattr(lemniscate, 'K'):
    criterion = NCECriterion(ndata)
else:
    criterion = nn.CrossEntropyLoss()

criterion_cld = nn.CrossEntropyLoss()
criterion_cld.to(device)

lemniscate.to(device)
criterion.to(device)

if args.test_only:
    acc = kNN(0, net, lemniscate, trainloader, testloader, 200, args.nce_t, 1)
    sys.exit(0)


# Training
torch.manual_seed(seed)
valdir = os.path.join('/data/carla_training_data/', 'val')
val_dataset = Dataset(valdir, n_frames, 1, gpu)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, 
                                         shuffle=True, num_workers=j)

ndata = train_dataset.__len__()
nce_k = 4096
nce_t = .07
nce_m = .5
iter_size = 1


#lemniscate = NCEAverage(gpu, low_dim, ndata, nce_k, nce_t, nce_m).to(get_device(gpu))
criterion = NCECriterion(ndata).to(get_device(gpu))

iter_size = 1

#hf = h5py.File('data_one.h5', 'w')
#with h5py.File('data.h5', "a") as f:
#    dset = f.create_dataset('img_names', (batch_size,), maxshape=(None,),
#                            dtype='S', chunks=(batch_size,))

print('create h5py data')
    
all_img_names = []
#all_imgs = []
all_id_nums = []
all_steer_truths = []
all_steer_preds = []
Exemplo n.º 8
0
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
           'ship', 'truck')
ndata = trainset.__len__()

print('==> Building model..')
net = models.__dict__['ResNet18'](low_dim=args.low_dim)

# define leminiscate
lemniscate = NCEAverage(args.low_dim,
                        ndata,
                        args.nce_t,
                        args.nce_m,
                        batchSize=args.b)
# define loss function
criterion = NCECriterion()

if device == 'cuda':
    net = torch.nn.DataParallel(net.cuda())
    lemniscate.cuda()
    criterion.cuda()
    cudnn.benchmark = True

# Model
if args.test_only or len(args.resume) > 0:
    # Load checkpoint.
    print('==> Resuming from checkpoint..')
    assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
    checkpoint = torch.load('./checkpoint/' + args.resume)
    net.load_state_dict(checkpoint['net'])
    lemniscate = checkpoint['lemniscate']
Exemplo n.º 9
0
    # get model and replace the original fc layer with your fc layer
    model_ft = models.resnet18(pretrained=True)
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, dataset_sizes)

    #model_ft.load_state_dict(torch.load(P_PATH))
    if use_gpu:
        model_ft = model_ft.cuda()

    # define loss function
    #criterion = nn.CrossEntropyLoss()
    #criterion = NEG_loss(dataset_sizes, dataset_sizes)

    lemniscate = NCEAverage(dataset_sizes, dataset_sizes, 100, 0.07, 0.5)
    criterion = NCECriterion(dataset_sizes)

    # Observe that all parameters are being optimized
    optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

    # Decay LR by a factor of 0.1 every 7 epochs
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft,
                                           step_size=7,
                                           gamma=0.1)

    model_ft = train_model(model=model_ft,
                           criterion=criterion,
                           optimizer=optimizer_ft,
                           scheduler=exp_lr_scheduler,
                           num_epochs=60)
    torch.save(model_ft.state_dict(), P_PATH)
Exemplo n.º 10
0
def main():

    global args, best_prec1
    args = parser.parse_args()

    # Initialize distributed processing
    args.distributed = args.world_size > 1
    if args.distributed:
        dist.init_process_group(backend=args.dist_backend,
                                init_method=args.dist_url,
                                world_size=args.world_size)

    # create model
    if args.pretrained:
        print("=> using pre-trained model '{}'".format(args.arch))
        model = models.__dict__[args.arch](pretrained=True,
                                           low_dim=args.low_dim)
    else:
        print("=> creating model '{}'".format(args.arch))
        model = models.__dict__[args.arch](low_dim=args.low_dim)

    if not args.distributed:
        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
            model.features = torch.nn.DataParallel(model.features)
            model.cuda()
        else:
            model = torch.nn.DataParallel(model).cuda()
    else:
        model.cuda()
        model = torch.nn.parallel.DistributedDataParallel(model)

    # Data loading code
    traindir = os.path.join(args.data, 'train')
    valdir = os.path.join(args.data, 'val')
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],  # ImageNet stats
        std=[0.229, 0.224, 0.225])
    #    normalize = transforms.Normalize(mean=[0.234, 0.191, 0.159],  # xView stats
    #                                     std=[0.173, 0.143, 0.127])

    print("Creating datasets")
    cj = args.color_jit
    train_dataset = datasets.ImageFolderInstance(
        traindir,
        transforms.Compose([
            transforms.Resize((224, 224)),
            #            transforms.Grayscale(3),
            #            transforms.ColorJitter(cj, cj, cj, cj), #transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
            transforms.RandomHorizontalFlip(),
            transforms.RandomVerticalFlip(),
            transforms.RandomRotation(45),
            transforms.ToTensor(),
            normalize,
        ]))

    if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_dataset)
    elif args.balanced_sampling:

        print("Using balanced sampling")
        # Here's where we compute the weights for WeightedRandomSampler
        class_counts = {v: 0 for v in train_dataset.class_to_idx.values()}
        for path, ndx in train_dataset.samples:
            class_counts[ndx] += 1
        total = float(np.sum([v for v in class_counts.values()]))
        class_probs = [
            class_counts[ndx] / total for ndx in range(len(class_counts))
        ]

        # make a list of class probabilities corresponding to the entries in train_dataset.samples
        reciprocal_weights = [
            class_probs[idx]
            for i, (_, idx) in enumerate(train_dataset.samples)
        ]

        # weights are the reciprocal of the above
        weights = (1 / torch.Tensor(reciprocal_weights))

        train_sampler = torch.utils.data.sampler.WeightedRandomSampler(
            weights, len(train_dataset), replacement=True)
    else:
        #if args.red_data is < 1, then the training is done with a subsamle of the total data. Otherwise it's the total data.
        data_size = len(train_dataset)
        sub_index = np.random.randint(0, data_size,
                                      round(args.red_data * data_size))
        sub_index.sort()
        train_sampler = torch.utils.data.sampler.SubsetRandomSampler(sub_index)

    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=args.batch_size,
                                               shuffle=(train_sampler is None),
                                               num_workers=args.workers,
                                               pin_memory=True,
                                               sampler=train_sampler)

    print("Training on", len(train_dataset.imgs),
          "images. Training batch size:", args.batch_size)

    if len(train_dataset.imgs) % args.batch_size != 0:
        print(
            "Warning: batch size doesn't divide the # of training images so ",
            len(train_dataset.imgs) % args.batch_size,
            "images will be skipped per epoch.")
        print("If you don't want to skip images, use a batch size in:",
              get_factors(len(train_dataset.imgs)))

    val_dataset = datasets.ImageFolderInstance(
        valdir,
        transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            normalize,
        ]))

    val_bs = [
        factor for factor in get_factors(len(val_dataset)) if factor < 500
    ][-1]
    val_bs = 100
    val_loader = torch.utils.data.DataLoader(val_dataset,
                                             batch_size=val_bs,
                                             shuffle=False,
                                             num_workers=args.workers,
                                             pin_memory=True)

    print("Validating on", len(val_dataset), "images. Validation batch size:",
          val_bs)

    # define lemniscate and loss function (criterion)
    ndata = train_dataset.__len__()
    if args.nce_k > 0:
        lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t,
                                args.nce_m)
        criterion = NCECriterion(ndata).cuda()
    else:
        lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t,
                                   args.nce_m).cuda()
        criterion = nn.CrossEntropyLoss().cuda()

    optimizer = torch.optim.SGD(model.parameters(),
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            model.load_state_dict(checkpoint['state_dict'])
            optimizer = FP16_Optimizer(optimizer,
                                       static_loss_scale=args.static_loss,
                                       verbose=False)
            optimizer.load_state_dict(checkpoint['optimizer'])
            args.start_epoch = checkpoint['epoch']
            #           best_prec1 = checkpoint['best_prec1']
            lemniscate = checkpoint['lemniscate']
            if args.select_load:
                pred = checkpoint['prediction']
            print("=> loaded checkpoint '{}' (epoch {}, best_prec1 )".format(
                args.resume,
                checkpoint['epoch']))  #, checkpoint['best_prec1']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    # optionally fine-tune a model trained on a different dataset
    elif args.fine_tune:
        print("=> loading checkpoint '{}'".format(args.fine_tune))
        checkpoint = torch.load(args.fine_tune)
        model.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        optimizer = FP16_Optimizer(optimizer,
                                   static_loss_scale=args.static_loss,
                                   verbose=False)
        print("=> loaded checkpoint '{}' (epoch {})".format(
            args.fine_tune, checkpoint['epoch']))
    else:
        optimizer = FP16_Optimizer(optimizer,
                                   static_loss_scale=args.static_loss,
                                   verbose=False)

    # Optionally recompute memory. If fine-tuning, then we must recompute memory
    if args.recompute_memory or args.fine_tune:

        # Aaron - Experiments show that iterating over torch.utils.data.DataLoader will skip the last few
        # unless the batch size evenly divides size of the data set. This shouldn't be the case
        # according to documentation, there's even a flag for drop_last, but it's not working

        # compute a good batch size for re-computing memory
        memory_bs = [
            factor for factor in get_factors(len(train_loader.dataset))
            if factor < 500
        ][-1]
        print("Recomputing memory using", train_dataset.root,
              "with a batch size of", memory_bs)
        transform_bak = train_loader.dataset.transform
        train_loader.dataset.transform = val_loader.dataset.transform
        temploader = torch.utils.data.DataLoader(
            train_loader.dataset,
            batch_size=memory_bs,
            shuffle=False,
            num_workers=train_loader.num_workers,
            pin_memory=True)
        lemniscate.memory = torch.zeros(len(train_loader.dataset),
                                        args.low_dim).cuda()
        model.eval()
        with torch.no_grad():
            for batch_idx, (inputs, targets,
                            indexes) in enumerate(tqdm.tqdm(temploader)):
                batchSize = inputs.size(0)
                features = model(inputs)
                lemniscate.memory[batch_idx * batchSize:batch_idx * batchSize +
                                  batchSize, :] = features.data
        train_loader.dataset.transform = transform_bak
        model.train()

    cudnn.benchmark = True

    if args.evaluate:
        kNN(model, lemniscate, train_loader, val_loader, args.K, args.nce_t)
        return

    begin_train_time = datetime.datetime.now()

    #    my_knn(model, lemniscate, train_loader, val_loader, args.K, args.nce_t, train_dataset, val_dataset)
    if args.tsne:
        labels = idx_to_name(train_dataset, args.graph_labels)
        tsne(lemniscate, args.tsne, labels)
    if args.pca:
        labels = idx_to_name(train_dataset, args.graph_labels)
        pca(lemniscate, labels)
    if args.view_knn:
        my_knn(model, lemniscate, train_loader, val_loader, args.K, args.nce_t,
               train_dataset, val_dataset)
    if args.kmeans:
        kmeans, yi = kmean(lemniscate, args.kmeans, 500, args.K, train_dataset)
        D, I = kmeans.index.search(lemniscate.memory.data.cpu().numpy(), 1)

        cent_group = {}
        data_cent = {}
        for n, i in enumerate(I):
            if i[0] not in cent_group.keys():
                cent_group[i[0]] = []
            cent_group[i[0]].append(n)
        data_cent[n] = i[0]

        train_sampler = torch.utils.data.sampler.SubsetRandomSampler(
            cent_group[0])
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=args.batch_size,
            shuffle=(train_sampler is None),
            num_workers=args.workers,
            pin_memory=True,
            sampler=train_sampler)

#        lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t, args.nce_m)
#        criterion = NCECriterion(ndata).cuda()

#    lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t, args.nce_m)

    if args.tsne_grid:
        tsne_grid(val_loader, model)
    if args.h_cluster:
        for size in range(2, 3):
            #        size = 20
            kmeans, topk = kmean(lemniscate, size, 500, 10, train_dataset)
            respred = torch.tensor([]).cuda()
            lab, idx = [[] for i in range(2)]
            num = 0
            '''
            for p,index,label in pred:
                respred = torch.cat((respred,p))
                if num == 0:
                    lab = label
                else:
                    lab += label
                idx.append(index)
                num+=1
            '''
            h_cluster(lemniscate, train_dataset, kmeans, topk,
                      size)  #, respred, lab, idx)

#    axis_explore(lemniscate, train_dataset)

#    kmeans_opt(lemniscate, 5)

    if args.select:
        if not args.select_load:
            pred = []

            if args.select_size:
                size = int(args.select_size * ndata)
            else:
                size = round(ndata / 100.0)

            sub_sample = np.random.randint(0, ndata, size=size)
            train_sampler = torch.utils.data.sampler.SubsetRandomSampler(
                sub_sample)
            train_loader = torch.utils.data.DataLoader(
                train_dataset,
                batch_size=args.batch_size,
                shuffle=(train_sampler is None),
                num_workers=args.workers,
                pin_memory=True,
                sampler=train_sampler)

            pred = div_train(train_loader, model, 0, pred)

        pred_features = []
        pred_labels = []
        pred_idx = []

        for inst in pred:
            feat, idx, lab = list(inst)
            pred_features.append(feat)
            pred_labels.append(lab)
            pred_idx.append(idx.data.cpu())

        if args.select_save:

            save_checkpoint(
                {
                    'epoch': args.start_epoch,
                    'arch': args.arch,
                    'state_dict': model.state_dict(),
                    'prediction': pred,
                    'lemniscate': lemniscate,
                    'optimizer': optimizer.state_dict(),
                }, 'select.pth.tar')

        min_idx = selection(pred_features, pred_idx, train_dataset,
                            args.select_num, args.select_thresh)

        train_sampler = torch.utils.data.sampler.SubsetRandomSampler(min_idx)

        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=args.batch_size,
            shuffle=(train_sampler is None),
            num_workers=args.workers,
            pin_memory=True,
            sampler=train_sampler)

        lemniscate = NCEAverage(args.low_dim, ndata, 20, args.nce_t,
                                args.nce_m)

        optimizer = torch.optim.SGD(model.parameters(),
                                    0.1,
                                    momentum=0.1,
                                    weight_decay=0.00001)

        optimizer = FP16_Optimizer(optimizer,
                                   static_loss_scale=args.static_loss,
                                   verbose=False)

        for epoch in range(50):
            if args.distributed:
                train_sampler.set_epoch(epoch)
            adjust_learning_rate(optimizer, epoch)

            if epoch % 1 == 0:
                save_checkpoint({
                    'epoch': epoch + 1,
                    'arch': args.arch,
                    'state_dict': model.state_dict(),
                    'lemniscate': lemniscate,
                    'optimizer': optimizer.state_dict(),
                })

            train(train_loader, model, lemniscate, criterion, optimizer, epoch)

        train_sampler = torch.utils.data.sampler.SubsetRandomSampler(sub_index)
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=args.batch_size,
            shuffle=(train_sampler is None),
            num_workers=args.workers,
            pin_memory=True,
            sampler=train_sampler)

        lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t,
                                args.nce_m)
        optimizer = torch.optim.SGD(model.parameters(),
                                    args.lr,
                                    momentum=args.momentum,
                                    weight_decay=args.weight_decay)
        optimizer = FP16_Optimizer(optimizer,
                                   static_loss_scale=args.static_loss,
                                   verbose=False)

    if args.kmeans_opt:
        kmeans_opt(lemniscate, 500)

    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            train_sampler.set_epoch(epoch)
        adjust_learning_rate(optimizer, epoch)

        if epoch % 1 == 0:
            # evaluate on validation set
            #prec1 = NN(epoch, model, lemniscate, train_loader, train_loader) # was evaluating on train
            #            prec1 = kNN(model, lemniscate, train_loader, val_loader, args.K, args.nce_t)
            # prec1 really should be renamed to prec5 as kNN now returns top5 score, but
            # it won't be backward's compatible as earlier models were saved with "best_prec1"

            # remember best prec@1 and save checkpoint
            #            is_best = prec1 > best_prec1
            #            best_prec1 = max(prec1, best_prec1)
            save_checkpoint({
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'lemniscate': lemniscate,
                #                'best_prec1': best_prec1,
                'optimizer': optimizer.state_dict(),
            })  # , is_best)

        # train for one epoch
        train(train_loader, model, lemniscate, criterion, optimizer, epoch)

#        kmeans,cent = kmeans()
#        group_train(train_loader, model, lemniscate, criterion, optimizer, epoch, kmeans, cent)

# print elapsed time
    end_train_time = datetime.datetime.now()
    d = end_train_time - begin_train_time
    print(
        "Trained for %d epochs. Elapsed time: %s days, %.2dh: %.2dm: %.2ds" %
        (len(range(args.start_epoch, args.epochs)), d.days, d.seconds // 3600,
         (d.seconds // 60) % 60, d.seconds % 60))
Exemplo n.º 11
0
def main():
    global args, best_prec1
    args = parser.parse_args()

    args.distributed = args.world_size > 1

    if args.distributed:
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                world_size=args.world_size)

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    train_transform = transforms.Compose(
        [
            transforms.RandomResizedCrop(224, scale=(0.3, 1.)),
            transforms.RandomGrayscale(p=0.5),
            transforms.ColorJitter(0.5, 0.5, 0.5, 0.5),
            transforms.RandomHorizontalFlip(),
            transforms.RandomVerticalFlip(),
            transforms.ToTensor(),
            normalize])

    val_transform = transforms.Compose(
        [
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            normalize])

    train_dataset = Foundation_Type_Binary(args.train_data, transform=train_transform, mask_buildings=args.mask_buildings)
    val_dataset = Foundation_Type_Binary(args.val_data, transform=val_transform, mask_buildings=args.mask_buildings)

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    val_loader = torch.utils.data.DataLoader(val_dataset,
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    model = resnet50(low_dim=args.low_dim)
    model = torch.nn.DataParallel(model).cuda()

    print ('Train dataset instances: {}'.format(len(train_loader.dataset)))
    print ('Val dataset instances: {}'.format(len(val_loader.dataset)))

    ndata = train_dataset.__len__()
    if args.nce_k > 0:
        lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t, args.nce_m).cuda()
        criterion = NCECriterion(ndata).cuda()
    else:
        lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t, args.nce_m).cuda()
        criterion = nn.CrossEntropyLoss().cuda()

    optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                 momentum=args.momentum,
                                 weight_decay=args.weight_decay)

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            args.epochs = args.start_epoch + args.epochs
            best_prec1 = checkpoint['best_prec1']

            missing_keys, unexpected_keys = model.load_state_dict(checkpoint['state_dict'], strict=False)
            if len(missing_keys) or len(unexpected_keys):
                print('Warning: Missing or unexpected keys found.')
                print('Missing: {}'.format(missing_keys))
                print('Unexpected: {}'.format(unexpected_keys))

            low_dim_checkpoint = checkpoint['lemniscate'].memory.shape[1]
            if low_dim_checkpoint == args.low_dim:
                lemniscate = checkpoint['lemniscate']
            else:
                print('Chosen low dim does not fit checkpoint. Assuming fine-tuning and not loading memory bank.')
            try:
                optimizer.load_state_dict(checkpoint['optimizer'])
            except ValueError:
                print('Training optimizer does not fit to checkpoint optimizer. Assuming fine-tuning and load optimizer from scratch. ')

            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    cudnn.benchmark = True

    if args.evaluate:
        kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
        return

    prec1 = NN(0, model, lemniscate, train_loader, val_loader)
    print('Start out precision: {}'.format(prec1))
    for epoch in range(args.start_epoch, args.epochs):

        adjust_learning_rate(optimizer, epoch)

        # train for one epoch
        train(train_loader, model, lemniscate, criterion, optimizer, epoch)

        # evaluate on validation set
        prec1 = NN(epoch, model, lemniscate, train_loader, val_loader)

        # remember best prec@1 and save checkpoint
        is_best = prec1 > best_prec1
        best_prec1 = max(prec1, best_prec1)
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'lemniscate': lemniscate,
            'best_prec1': best_prec1,
            'optimizer' : optimizer.state_dict(),
        }, is_best, args.name)
Exemplo n.º 12
0
cudnn.benchmark = True

# Model
if args.test_only or len(args.resume) > 0:
    # Load checkpoint.
    print('==> Resuming from checkpoint..')
    assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
    checkpoint = torch.load('./checkpoint/' + args.resume)
    net.load_state_dict(checkpoint['net'])
    lemniscate = checkpoint['lemniscate']
    best_acc = checkpoint['acc']
    start_epoch = checkpoint['epoch']

# define loss function
if hasattr(lemniscate, 'K'):
    criterion = NCECriterion(ndata)
else:
    criterion = nn.CrossEntropyLoss()

net.cuda()
lemniscate.cuda()
criterion.cuda()

if args.test_only:
    acc = kNN(0, net, lemniscate, trainloader, testloader, 200, args.nce_t, 1)
    sys.exit(0)

optimizer = optim.SGD(net.parameters(),
                      lr=args.lr,
                      momentum=0.9,
                      weight_decay=5e-4)
Exemplo n.º 13
0
    net = checkpoint['net']
    lemniscate = checkpoint['lemniscate']
    best_acc = checkpoint['acc']
    start_epoch = checkpoint['epoch']
else:
    print('==> Building model..')
    net = models.__dict__['ResNet18'](low_dim=args.low_dim)
    # define leminiscate
    if args.nce_k > 0:
        lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t, args.nce_m)
    else:
        lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t, args.nce_m)

# define loss function
if hasattr(lemniscate, 'K'):
    criterion = NCECriterion(ndata)
else:
    criterion = nn.CrossEntropyLoss()

if use_cuda:
    net.cuda()
    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
    lemniscate.cuda()
    criterion.cuda()
    cudnn.benchmark = True

if args.test_only:
    acc = kNN(0, net, lemniscate, trainloader, testloader, 200, args.nce_t, 1)
    sys.exit(0)

optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
Exemplo n.º 14
0
def build_model():
    best_acc = 0  # best test accuracy
    start_epoch = 0  # start from epoch 0 or last checkpoint epoch

    if args.architecture == 'resnet18':
        net = models.__dict__['resnet18_cifar'](low_dim=args.low_dim)
    elif args.architecture == 'wrn-28-2':
        net = models.WideResNet(depth=28,
                                num_classes=args.low_dim,
                                widen_factor=2,
                                dropRate=0).to(args.device)
    elif args.architecture == 'wrn-28-10':
        net = models.WideResNet(depth=28,
                                num_classes=args.low_dim,
                                widen_factor=10,
                                dropRate=0).to(args.device)

    # define leminiscate
    if args.nce_k > 0:
        lemniscate = NCEAverage(args.low_dim, args.ndata, args.nce_k,
                                args.nce_t, args.nce_m)
    else:
        lemniscate = LinearAverage(args.low_dim, args.ndata, args.nce_t,
                                   args.nce_m)

    if args.device == 'cuda':
        net = torch.nn.DataParallel(net,
                                    device_ids=range(
                                        torch.cuda.device_count()))
        cudnn.benchmark = True

    optimizer = optim.SGD(net.parameters(),
                          lr=args.lr,
                          momentum=0.9,
                          weight_decay=args.weight_decay,
                          nesterov=True)
    # Model
    if args.test_only or len(args.resume) > 0:
        # Load checkpoint.
        print('==> Resuming from checkpoint..')
        checkpoint = torch.load(args.resume)
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lemniscate = checkpoint['lemniscate']
        best_acc = checkpoint['acc']
        start_epoch = checkpoint['epoch'] + 1

    if args.lr_scheduler == 'multi-step':
        if args.epochs == 200:
            steps = [60, 120, 160]
        elif args.epochs == 600:
            steps = [180, 360, 480, 560]
        else:
            raise RuntimeError(
                f"need to config steps for epoch = {args.epochs} first.")
        scheduler = lr_scheduler.MultiStepLR(optimizer,
                                             steps,
                                             gamma=0.2,
                                             last_epoch=start_epoch - 1)
    elif args.lr_scheduler == 'cosine':
        scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
                                                   args.epochs,
                                                   eta_min=0.00001,
                                                   last_epoch=start_epoch - 1)
    elif args.lr_scheduler == 'cosine-with-restart':
        scheduler = CosineAnnealingLRWithRestart(optimizer,
                                                 eta_min=0.00001,
                                                 last_epoch=start_epoch - 1)
    else:
        raise ValueError("not supported")

    # define loss function
    if hasattr(lemniscate, 'K'):
        criterion = NCECriterion(args.ndata)
    else:
        criterion = nn.CrossEntropyLoss()

    net.to(args.device)
    lemniscate.to(args.device)
    criterion.to(args.device)

    return net, lemniscate, optimizer, criterion, scheduler, best_acc, start_epoch
Exemplo n.º 15
0
def main():
    global args, best_prec1, min_avgloss
    args = parser.parse_args()
    input("Begin the {} time's training".format(args.train_time))
    writer_log_dir = "/data/fhz/unsupervised_recommendation/idfe_runs/idfe_train_time:{}".format(
        args.train_time)
    writer = SummaryWriter(log_dir=writer_log_dir)
    if args.dataset == "lung":
        # build dataloader,val_dloader will be build in test function
        model = idfe.IdFe3d(feature_dim=args.latent_dim)
        model.encoder = torch.nn.DataParallel(model.encoder)
        model.linear_map = torch.nn.DataParallel(model.linear_map)
        model = model.cuda()
        train_datalist, test_datalist = multi_cross_validation()
        ndata = len(train_datalist)
    elif args.dataset == "gland":
        dataset_path = "/data/fhz/MICCAI2015/npy"
        model = idfe.IdFe2d(feature_dim=args.latent_dim)
        model.encoder = torch.nn.DataParallel(model.encoder)
        model.linear_map = torch.nn.DataParallel(model.linear_map)
        model = model.cuda()
        train_datalist = glob(path.join(dataset_path, "train", "*.npy"))
        ndata = len(train_datalist)
    else:
        raise FileNotFoundError("Dataset {} Not Found".format(args.dataset))
    if args.nce_k > 0:
        """
        Here we use NCE to calculate loss
        """
        lemniscate = NCEAverage(args.latent_dim, ndata, args.nce_k, args.nce_t,
                                args.nce_m).cuda()
        criterion = NCECriterion(ndata).cuda()
    else:
        lemniscate = LinearAverage(args.latent_dim, ndata, args.nce_t,
                                   args.nce_m).cuda()
        criterion = nn.CrossEntropyLoss().cuda()
    optimizer = torch.optim.SGD(model.parameters(),
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            min_avgloss = checkpoint['min_avgloss']
            model.encoder.load_state_dict(checkpoint['encoder_state_dict'])
            model.linear_map.load_state_dict(
                checkpoint['linear_map_state_dict'])
            lemniscate = checkpoint['lemniscate']
            optimizer.load_state_dict(checkpoint['optimizer'])
            train_datalist = checkpoint['train_datalist']
            print("=> loaded checkpoint '{}' (epoch {})".format(
                args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))
    if args.dataset == "lung":
        train_dset = LungDataSet(data_path_list=train_datalist,
                                 augment_prob=args.aug_prob)
        train_dloader = DataLoader(dataset=train_dset,
                                   batch_size=args.batch_size,
                                   shuffle=True,
                                   num_workers=args.workers,
                                   pin_memory=True)
    elif args.dataset == "gland":
        train_dset = GlandDataset(data_path_list=train_datalist,
                                  need_seg_label=False,
                                  augment_prob=args.aug_prob)
        train_dloader = DataLoader(dataset=train_dset,
                                   batch_size=args.batch_size,
                                   shuffle=True,
                                   num_workers=args.workers,
                                   pin_memory=True)
    else:
        raise FileNotFoundError("Dataset {} Not Found".format(args.dataset))
    for epoch in range(args.start_epoch, args.epochs):
        adjust_learning_rate(optimizer, epoch)
        epoch_loss = train(train_dloader,
                           model=model,
                           lemniscate=lemniscate,
                           criterion=criterion,
                           optimizer=optimizer,
                           epoch=epoch,
                           writer=writer,
                           dataset=args.dataset)
        if (epoch + 1) % 5 == 0:
            if args.dataset == "lung":
                """
                Here we define the best point as the minimum average epoch loss
                
                """
                accuracy = list([])
                # for i in range(5):
                #     train_feature = lemniscate.memory.clone()
                #     test_datalist = train_datalist[five_cross_idx[i]:five_cross_idx[i + 1]]
                #     test_feature = train_feature[five_cross_idx[i]:five_cross_idx[i + 1], :]
                #     train_indices = [train_datalist.index(d) for d in train_datalist if d not in test_datalist]
                #     tmp_train_feature = torch.index_select(train_feature, 0, torch.tensor(train_indices).cuda())
                #     tmp_train_datalist = [train_datalist[i] for i in train_indices]
                #     test_label = np.array(
                #         [int(eval(re.match("(.*)_(.*)_annotations.npy", path.basename(raw_cube_path)).group(2)) > 3)
                #          for raw_cube_path in test_datalist], dtype=np.float)
                #     tmp_train_label = np.array(
                #         [int(eval(re.match("(.*)_(.*)_annotations.npy", path.basename(raw_cube_path)).group(2)) > 3)
                #          for raw_cube_path in tmp_train_datalist], dtype=np.float)
                #     accuracy.append(
                #         kNN(tmp_train_feature, tmp_train_label, test_feature, test_label, K=20, sigma=1 / 10))
                # accuracy = mean(accuracy)
                is_best = (epoch_loss < min_avgloss)
                min_avgloss = min(epoch_loss, min_avgloss)
                save_checkpoint(
                    {
                        'epoch': epoch + 1,
                        "train_time": args.train_time,
                        "encoder_state_dict": model.encoder.state_dict(),
                        "linear_map_state_dict": model.linear_map.state_dict(),
                        'lemniscate': lemniscate,
                        'min_avgloss': min_avgloss,
                        'dataset': args.dataset,
                        'optimizer': optimizer.state_dict(),
                        'train_datalist': train_datalist
                    }, is_best)
                # knn_text = "In epoch :{} the five cross validation accuracy is :{}".format(epoch, accuracy * 100.0)
                # # print(knn_text)
                # writer.add_text("knn/text", knn_text, epoch)
                # writer.add_scalar("knn/accuracy", accuracy, global_step=epoch)
            elif args.dataset == "gland":
                is_best = (epoch_loss < min_avgloss)
                min_avgloss = min(epoch_loss, min_avgloss)
                save_checkpoint(
                    {
                        'epoch': epoch + 1,
                        "train_time": args.train_time,
                        "encoder_state_dict": model.encoder.state_dict(),
                        "linear_map_state_dict": model.linear_map.state_dict(),
                        'lemniscate': lemniscate,
                        'min_avgloss': min_avgloss,
                        'dataset': args.dataset,
                        'optimizer': optimizer.state_dict(),
                        'train_datalist': train_datalist,
                    }, is_best)