Function train

methods/representation_learning/representation_learning.py:179–430

(projectors, student, train_loaders, test_loader, unlabelled_train_loader, whole_train_test_loader, args)

def train(projectors, student, train_loaders, test_loader, unlabelled_train_loader, whole_train_test_loader, args):

    optimizer = SGD(list(projectors.parameters()) + list(student.parameters()), lr=args.lr, momentum=args.momentum,
                    weight_decay=args.weight_decay)

    exp_lr_scheduler = lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=args.epochs,
            eta_min=args.lr * 1e-3,
        )

    sup_con_crit = SupConLoss()
    best_test_acc_lab = 0
    best_acc_lab = 0

    for epoch in range(args.epochs):

        loss_record = AverageMeter()
        train_acc_record = AverageMeter()

        student.train()
        projectors.train()

        loaders =[iter(l) for l in train_loaders]

        max_len  = max([len(loader) for loader in loaders])

        # Load data from each dataloaders, the total number of dataloader is expert_num + 1
        for _ in tqdm(range(max_len)):
            fine_loss = 0
            expert_images = []
            expert_class_labels = []
            expert_mask_lab = []
            
            for idx, loader in enumerate(loaders):
                try:
                    item = next(loader)
                except StopIteration:
                    loaders[idx] = iter(train_loaders[idx])
                    item = next(loaders[idx])

                images, class_labels, uq_idxs, mask_lab = item
                mask_lab = mask_lab[:, 0]
                images = torch.cat(images, dim=0)   # [B*2/num_experts,3,224,224]
                if args.use_global_con:
                    # Load subset data
                    if idx < args.experts_num:
                        expert_images.append(images)
                        expert_class_labels.append(class_labels)
                        expert_mask_lab.append(mask_lab.bool())

                    # Load whole dataset
                    else:
                        all_images = images.to(device)
                        all_class_labels = class_labels.to(device)
                        all_mask_lab = mask_lab.to(device).bool()
                else:
                    expert_images.append(images)