def run_simclr(args, model, log_dir): args = copy.deepcopy(args) del args.la for k in args.simclr: args[k] = args.simclr[k] assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2." dataset = ContrastiveLearningDataset(args.dataset_path) train_dataset = dataset.get_dataset(args.dataset_name, args.n_views) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) log_dir = os.path.join(log_dir, 'simclr') os.makedirs(log_dir) # It’s a no-op if the 'gpu_index' argument is a negative integer or None. with torch.cuda.device(args.gpu_index): simclr = SimCLR(model=model, optimizer=optimizer, scheduler=scheduler, args=args, log_dir=log_dir) simclr.train(train_loader)
def main(): args = parser.parse_args() assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2." # check if gpu training is available if not args.disable_cuda and torch.cuda.is_available(): args.device = torch.device('cuda') cudnn.deterministic = True cudnn.benchmark = True else: args.device = torch.device('cpu') args.gpu_index = -1 dataset = ContrastiveLearningDataset(args.data) train_dataset = dataset.get_dataset(args.dataset_name, args.n_views) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim) optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) # It’s a no-op if the 'gpu_index' argument is a negative integer or None. with torch.cuda.device(args.gpu_index): simclr = SimCLR(model=model, optimizer=optimizer, scheduler=scheduler, args=args) simclr.train(train_loader)
class ContrastiveLearning(LightningModule): def __init__(self, args, encoder): super().__init__() self.save_hyperparameters(args) self.encoder = encoder self.n_features = (self.encoder.fc.in_features ) # get dimensions of last fully-connected layer self.model = SimCLR(self.encoder, self.hparams.projection_dim, self.n_features) self.criterion = self.configure_criterion() def forward(self, x_i, x_j): _, _, z_i, z_j = self.model(x_i, x_j) loss = self.criterion(z_i, z_j) return loss def training_step(self, batch, batch_idx): x, _ = batch x_i = x[:, 0, :].unsqueeze(dim=1) x_j = x[:, 1, :].unsqueeze(dim=1) loss = self.forward(x_i, x_j) self.log("Train/loss", loss) return loss def configure_criterion(self): # PT lightning aggregates differently in DP mode if self.hparams.accelerator == "dp" and self.hparams.gpus: batch_size = int(self.hparams.batch_size / self.hparams.gpus) else: batch_size = self.hparams.batch_size criterion = NT_Xent(batch_size, self.hparams.temperature, world_size=1) return criterion def configure_optimizers(self): scheduler = None if self.hparams.optimizer == "Adam": optimizer = torch.optim.Adam(self.model.parameters(), lr=3e-4) elif self.hparams.optimizer == "LARS": # optimized using LARS with linear learning rate scaling # (i.e. LearningRate = 0.3 × BatchSize/256) and weight decay of 10−6. learning_rate = 0.3 * self.hparams.batch_size / 256 optimizer = LARS( self.model.parameters(), lr=learning_rate, weight_decay=self.hparams.weight_decay, exclude_from_weight_decay=["batch_normalization", "bias"], ) # "decay the learning rate with the cosine decay schedule without restarts" scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, self.hparams.max_epochs, eta_min=0, last_epoch=-1) else: raise NotImplementedError if scheduler: return {"optimizer": optimizer, "lr_scheduler": scheduler} else: return {"optimizer": optimizer}
def main(): # Totalcases = 1051 # US-4 # Totalcases = 63 # CLUST # Totalcases = 296 # Liver Totalcases = 22 # Butterfly # Totalcases = 670 # COVID19 # LabelRate = [0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1] LabelRate = [1] # DataRate = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] DataRate = [1] Checkpoint_Num = 1 for LL in range(len(LabelRate)): for DD in range(len(DataRate)): LabelList = random.sample(range(0, Totalcases), math.ceil(Totalcases * LabelRate[LL])) DataList = random.sample(range(0, Totalcases), math.ceil(Totalcases * DataRate[DD])) config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader) dataset = DataSetWrapper(config['batch_size'], LabelList, DataList, Checkpoint_Num, **config['dataset']) lumbda = [0.1, 0.2, 0.3, 0.4, 0.5] # Semi-supervised CL, default=0.2 for i in range(len(lumbda)): simclr = SimCLR(dataset, config, lumbda[i], Checkpoint_Num) simclr.train() Checkpoint_Num += 1
def main(): config = yaml.load(open("config-cifar-eval.yaml", "r"), Loader=yaml.FullLoader) dataset = DataSetWrapper(config['batch_size'], **config['dataset']) simclr = SimCLR(dataset, config) simclr.eval_frozen()
def __init__(self): super().__init__() # initialize ResNet self.encoder = ENCODER self.n_features = ENCODER.fc.in_features # get dimensions of fc layer self.model = SimCLR(ENCODER, PROJ_DIM, self.n_features) self.criterion = NT_Xent(BATCH_SIZE, TEMPERATURE, world_size=1)
def __init__(self, net): super().__init__() # initialize ResNet self.encoder = net self.n_features = net.fc.in_features # get dimensions of fc layer self.model = SimCLR(net, 256, self.n_features) self.criterion = NT_Xent(32, 0.5, world_size=1)
def main(): args = parser.parse_args() assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2." # check if gpu training is available if not args.disable_cuda and torch.cuda.is_available(): args.device = torch.device('cuda') cudnn.deterministic = True cudnn.benchmark = True else: args.device = torch.device('cpu') args.gpu_index = -1 dataset = ContrastiveLearningDataset(args.data) train_dataset = dataset.get_dataset(args.dataset_name, args.n_views) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim).to(args.device) optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) if args.resume: if os.path.isfile(args.resume): print("=> loading resumed checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume, map_location=torch.device('cpu')) args.start_epoch = checkpoint['epoch'] args.start_iter = checkpoint['iteration'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) print("=> loaded resumed checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("[Warning] no checkpoint found at '{}'".format(args.resume)) else: args.start_iter = 0 args.start_epoch = 0 # It’s a no-op if the 'gpu_index' argument is a negative integer or None. with torch.cuda.device(args.gpu_index): simclr = SimCLR(model=model, optimizer=optimizer, scheduler=scheduler, args=args) simclr.train(train_loader)
def main(save_loc: str = None): config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader) dataset = DataSetWrapper(config['batch_size'], **config['dataset']) simclr = SimCLR(dataset, config) callback = get_valid_callback(save_loc) simclr.train(callback)
def __init__(self, args, encoder): super().__init__() self.save_hyperparameters(args) self.encoder = encoder self.n_features = (self.encoder.fc.in_features ) # get dimensions of last fully-connected layer self.model = SimCLR(self.encoder, self.hparams.projection_dim, self.n_features) self.criterion = self.configure_criterion()
def main(): parser = argparse.ArgumentParser() parser.add_argument('--magnification', type=str, default='20x') args = parser.parse_args() config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader) dataset = DataSetWrapper(config['batch_size'], **config['dataset']) generate_csv(args.magnification) simclr = SimCLR(dataset, config) simclr.train()
class ContrastiveLearning(LightningModule): def __init__(self, args): super().__init__() self.hparams = args # initialize ResNet self.encoder = get_resnet(self.hparams.resnet, pretrained=False) self.n_features = self.encoder.fc.in_features # get dimensions of fc layer self.model = SimCLR(self.encoder, self.hparams.projection_dim, self.n_features) self.criterion = NT_Xent(self.hparams.batch_size, self.hparams.temperature, world_size=1) def forward(self, x_i, x_j): h_i, h_j, z_i, z_j = self.model(x_i, x_j) loss = self.criterion(z_i, z_j) return loss def training_step(self, batch, batch_idx): # training_step defined the train loop. It is independent of forward (x_i, x_j), _ = batch loss = self.forward(x_i, x_j) return loss def configure_criterion(self): criterion = NT_Xent(self.hparams.batch_size, self.hparams.temperature) return criterion def configure_optimizers(self): scheduler = None if self.hparams.optimizer == "Adam": optimizer = torch.optim.Adam(self.model.parameters(), lr=3e-4) elif self.hparams.optimizer == "LARS": # optimized using LARS with linear learning rate scaling # (i.e. LearningRate = 0.3 × BatchSize/256) and weight decay of 10−6. learning_rate = 0.3 * args.batch_size / 256 optimizer = LARS( self.model.parameters(), lr=learning_rate, weight_decay=args.weight_decay, exclude_from_weight_decay=["batch_normalization", "bias"], ) # "decay the learning rate with the cosine decay schedule without restarts" scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, args.epochs, eta_min=0, last_epoch=-1) else: raise NotImplementedError if scheduler: return {"optimizer": optimizer, "lr_scheduler": scheduler} else: return {"optimizer": optimizer}
def main(): config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader) print(config) dataset = DataSetWrapper(config['batch_size'], **config['dataset']) simclr = SimCLR(dataset, config) simclr.train() device = 'cuda' if torch.cuda.is_available() else 'cpu' print(device) model = simclr.model eval(model, './data/', device, config)
def main(): config = yaml.load( open("config.yaml", "r"), Loader=yaml.FullLoader ) # This is a config file that consist of all the paramters neeed config the data loader. dataset = DataSetWrapper( config['batch_size'], **config['dataset'] ) #This is the dataloader object in pytorch. Refer to data_aug/dataset_wrapper.py simclr = SimCLR(dataset, config) #the model and stuff simclr.train()
def main(): parser = argparse.ArgumentParser() parser.add_argument('config', type=str, help='Path to config') config_path = parser.parse_args().config exp_name = os.path.splitext(os.path.basename(config_path))[0] config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) dataset = DataSetWrapper(batch_size=config['batch_size'], **config['dataset']) simclr = SimCLR(dataset, config, config_path, exp_name) simclr.train()
def main(): parser = argparse.ArgumentParser() parser.add_argument('--level', type=str, default='low', help='Magnification level to compute embedder (low/high)') parser.add_argument('--multiscale', type=int, default=0, help='Whether the patches are cropped from multiscale (0/1-no/yes)') parser.add_argument('--dataset', type=str, default='TCGA-lung', help='Dataset folder name') args = parser.parse_args() config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader) gpu_ids = eval(config['gpu_ids']) os.environ['CUDA_VISIBLE_DEVICES']=','.join(str(x) for x in gpu_ids) dataset = DataSetWrapper(config['batch_size'], **config['dataset']) generate_csv(args) simclr = SimCLR(dataset, config) simclr.train()
def __init__(self, args): super().__init__() self.hparams = args # initialize ResNet self.encoder = get_resnet(self.hparams.resnet, pretrained=False) self.n_features = self.encoder.fc.in_features # get dimensions of fc layer self.model = SimCLR(self.encoder, self.hparams.projection_dim, self.n_features) self.criterion = NT_Xent(self.hparams.batch_size, self.hparams.temperature, world_size=1)
def main(): config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader) os.environ['CUDA_VISIBLE_DEVICES'] = config['gpu']['gpu_ids'] print(os.environ['CUDA_VISIBLE_DEVICES']) batch_size = config['train']['train_batch_size_per_gpu'] * config['gpu'][ 'gpunum'] #if torch.cuda.is_available(): # batch_size = batch_size * config['gpu']['gpunum'] dataset = DataSetWrapper(batch_size, **config['dataset']) simclr = SimCLR(dataset, config) simclr.train()
class ContrastiveLearning(pl.LightningModule): def __init__(self): super().__init__() # initialize ResNet self.encoder = ENCODER self.n_features = ENCODER.fc.in_features # get dimensions of fc layer self.model = SimCLR(ENCODER, PROJ_DIM, self.n_features) self.criterion = NT_Xent(BATCH_SIZE, TEMPERATURE, world_size=1) def forward(self, x_i, x_j): h_i, h_j, z_i, z_j = self.model(x_i, x_j) loss = self.criterion(z_i, z_j) return loss def training_step(self, batch, batch_idx): (x_i, x_j) = batch loss = self.forward(x_i, x_j) self.log('loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True) return loss def configure_criterion(self): criterion = NT_Xent(BATCH_SIZE, self.hparams.temperature) return criterion def configure_optimizers(self): optimizer = torch.optim.Adam(self.model.parameters(), lr=LR) return optimizer
class ContrastiveLearning(pl.LightningModule): def __init__(self, net): super().__init__() # initialize ResNet self.encoder = net self.n_features = net.fc.in_features # get dimensions of fc layer self.model = SimCLR(net, 256, self.n_features) self.criterion = NT_Xent(32, 0.5, world_size=1) def forward(self, x_i, x_j): h_i, h_j, z_i, z_j = self.model(x_i, x_j) loss = self.criterion(z_i, z_j) return loss def training_step(self, batch, batch_idx): (x_i, x_j) = batch loss = self.forward(x_i, x_j) return loss def configure_criterion(self): criterion = NT_Xent(32, self.hparams.temperature) return criterion def configure_optimizers(self): optimizer = torch.optim.Adam(self.model.parameters()) return optimizer
def __init__(self, args): super().__init__() # self.hparams = args self.args = args # initialize ResNet self.encoder = get_resnet(self.args.resnet, pretrained=self.args.pretrain) self.n_features = self.encoder.fc.in_features # get dimensions of fc layer self.model = SimCLR(self.encoder, self.args.h_dim, self.args.projection_dim, self.n_features, self.args.n_classes) self.test_outputs = np.array([]) self.criterion = NT_Xent(self.args.batch_size, self.args.temperature, world_size=1)
def load_model(args): model = SimCLR(backbone=args.backbone, projection_dim=args.projection_dim, pretrained=args.pretrained, normalize=args.normalize) if args.inference: model.load_state_dict( torch.load("SimCLR_{}_epoch90.pth".format(args.backbone))) model = model.to(args.device) scheduler = None if args.optimizer == "Adam": optimizer = Adam(model.parameters(), lr=3e-4) # TODO: LARS elif args.optimizer == "LARS": # optimized using LARS with linear learning rate scaling # (i.e. LearningRate = 0.3 × BatchSize/256) and weight decay of 10−6. learning_rate = 0.3 * args.batch_size / 256 optimizer = LARS( model.parameters(), lr=learning_rate, weight_decay=args.weight_decay, exclude_from_weight_decay=["batch_normalization", "bias"], ) # "decay the learning rate with the cosine decay schedule without restarts" scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs, eta_min=0, last_epoch=-1) else: raise NotImplementedError return model, optimizer, scheduler
def build_model(cfg): cfg.build_backbone = build_backbone model = SimCLR(cfg) logger = logging.getLogger(__name__) logger.info("Model:\n{}".format(model)) return model
def main(): parser=argparse.ArgumentParser() parser.add_argument('--config_path', type=str, default="config.yaml", help='Name of config file to use') config = parser.parse_args() print(f"Loading {config.config_path}") config = yaml.load(open(config.config_path, "r"), Loader=yaml.FullLoader) print(f"Model in this file is {config['model']['base_model']}") if torch.cuda.is_available() and config['allow_multiple_gpu']: gpu_count = torch.cuda.device_count() if gpu_count > 1: print(f'There are {gpu_count} GPUs with the current setup, so we will increase batch size and later run the model on all GPUs') config['batch_size'] *= gpu_count else: print("There is only 1 GPU available") else: print("There are no GPUs available") dataset = DataSetWrapper(config['batch_size'], **config['dataset']) simclr = SimCLR(dataset, config) simclr.train()
def main(): args = update_parser_args(task="pre_train") CL_dataset = ContrastiveLearningDataset(args.data) train_loader = torch.utils.data.DataLoader(dataset=CL_dataset.get_dataset( args.dataset_name, args.n_views), batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) simclr_model = PreTrain_NN(base_model=args.arch, out_dim=args.out_dim) optimizer = torch.optim.Adam(simclr_model.parameters(), args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) with torch.cuda.device( args.gpu_index ): # It’s a no-op if the 'gpu_index' argument is a negative integer or None. May need to update this for multiGPU? simclr_pt = SimCLR(model=simclr_model, optimizer=optimizer, scheduler=scheduler, args=args) simclr_pt.train(train_loader)
def main(): config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader) dataset = DataSetWrapper(config['batch_size'], **config['dataset']) if config["exp_type"] == "adversarial": print("Use Adversarial Augmentation.") simclr = SimCLRAdv(dataset, config) elif config["exp_type"] == "icm_adversarial": simclr = IcmSimCLR(dataset, config) elif config["exp_type"] == "icm_adversarial_v2": simclr = IcmSimCLRv2(dataset, config) elif config["exp_type"] == "icm_adversarial_v3": simclr = IcmSimCLRv3(dataset, config) elif config["exp_type"] == "normal": simclr = SimCLR(dataset, config) else: raise ValueError("Unrecognized experiment type: {}".format( config["exp_type"])) simclr.train()
def main(): config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader) dataset = DataSetWrapper(config['batch_size'], **config['dataset']) simclr = SimCLR(dataset, config) if config['mode'] == 'train': simclr.train() if config['mode'] == 'test': simclr.test() if config['mode'] == 'eval': logistic = ResNetFeatureExtractor(config) X_train_feature, y_train, X_test_feature, y_test = logistic.get_resnet_features( ) classifier = LogiticRegressionEvaluator(2048, 10) classifier.train(X_train_feature, y_train, X_test_feature, y_test)
num_workers=args.workers, ) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=args.logistic_batch_size, shuffle=False, drop_last=True, num_workers=args.workers, ) encoder = get_resnet(args.resnet, pretrained=False) n_features = encoder.fc.in_features # get dimensions of fc layer # load pre-trained model from checkpoint simclr_model = SimCLR(encoder, args.projection_dim, n_features) model_fp = os.path.join(args.model_path, "checkpoint_{}.tar".format(args.epoch_num)) simclr_model.load_state_dict(torch.load(model_fp, map_location=args.device.type)) simclr_model = simclr_model.to(args.device) simclr_model.eval() ## Logistic Regression n_classes = 10 # CIFAR-10 / STL-10 model = LogisticRegression(simclr_model.n_features, n_classes) model = model.to(args.device) optimizer = torch.optim.Adam(model.parameters(), lr=3e-4) criterion = torch.nn.CrossEntropyLoss() print("### Creating features from pre-trained context model ###") (train_X, train_y, test_X, test_y) = get_features(
def __init__(self, input_dim=64, output_dim=8, **kwargs): super().__init__() self.save_hyperparameters() self.model = LeadWishCNN(num_cls=output_dim, dim=input_dim) if self.hparams.ptmodel == "moco2": self.ptmodel = MocoV2.load_from_checkpoint( self.hparams.ptmodel_path, strict=False) self.ptmodel = self.ptmodel.encoder_q self.ptmodel = torch.nn.Sequential( *(list(self.ptmodel.children())[:-1])) self.model.feature = self.ptmodel if self.hparams.ptmodel == "simclr": self.ptmodel = SimCLR.load_from_checkpoint( self.hparams.ptmodel_path, strict=False) self.model.feature = self.ptmodel.encoder.feature if self.hparams.ptmodel == "byol": self.ptmodel = BYOL.load_from_checkpoint(self.hparams.ptmodel_path, strict=False) self.model.feature = self.ptmodel.online_network.encoder if self.hparams.ptmodel == "simsiam": self.ptmodel = SimSiam.load_from_checkpoint( self.hparams.ptmodel_path, strict=False) self.model.feature = self.ptmodel.online_network.encoder.feature if self.hparams.ptmodel == "swav": self.ptmodel = SwAV.load_from_checkpoint(self.hparams.ptmodel_path, strict=False) self.model.feature = self.ptmodel.model.feature self.model.feature = nn.Sequential(self.model.feature, nn.AdaptiveAvgPool2d(1)) if self.hparams.ptmodel == "imagenet": self.ptmodel = torchvision.models.resnet18(pretrained=True) self.ptmodel = torch.nn.Sequential( *(list(self.ptmodel.children())[:-1])) self.ptmodel[0] = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) self.model.feature = self.ptmodel if self.hparams.ptmodel == "cpc": self.ptmodel = CPCV2.load_from_checkpoint( self.hparams.ptmodel_path) self.model.feature = self.ptmodel.encoder # frozen conv layer if self.hparams.frozen == "true": print("Frozen all convolution layers") for m in self.model.modules(): if isinstance(m, torch.nn.Conv2d): m.weight.requires_grad = False elif self.hparams.frozen == "false": print("Train all layers") elif self.hparams.frozen == "group": print("Hierarchical learning rate") if self.hparams.classifier == "linear": if not self.hparams.dropout: self.model.classifier = nn.Linear(input_dim * 8 * 12, output_dim) elif self.hparams.classifier == "mlp": hidden_dim = 512 self.model.classifier == nn.Sequential( nn.Dropout(0.5), nn.Linear(input_dim * 8 * 12, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(inplace=True), nn.Dropout(0.5), nn.Linear(hidden_dim, output_dim), ) # LeadWishCNN 分组 self.model_groups = [] model_list = list(self.model.feature.children()) self.model_groups.append(nn.Sequential(*model_list[:5])) self.model_groups.append(nn.Sequential(*model_list[5:6])) self.model_groups.append(nn.Sequential(*model_list[6:7])) self.model_groups.append(nn.Sequential(*model_list[7:])) self.model_groups.append(self.model.classifier)
class ContrastiveLearning(LightningModule): def __init__(self, args): super().__init__() # self.hparams = args self.args = args # initialize ResNet self.encoder = get_resnet(self.args.resnet, pretrained=self.args.pretrain) self.n_features = self.encoder.fc.in_features # get dimensions of fc layer self.model = SimCLR(self.encoder, self.args.h_dim, self.args.projection_dim, self.n_features, self.args.n_classes) self.test_outputs = np.array([]) self.criterion = NT_Xent(self.args.batch_size, self.args.temperature, world_size=1) def forward(self, x_i, x_j): h_i, h_j, z_i, z_j = self.model(x_i, x_j) return h_i, h_j, z_i, z_j def training_step(self, batch, batch_idx): # training_step defined the train loop. It is independent of forward if self.args.dataset == "VGG_Face": (x_i, x_j) = batch else: (x_i, x_j), _ = batch h_i, h_j, z_i, z_j = self.forward(x_i, x_j) loss = self.criterion(z_i, z_j) self.log("Training Loss", loss, on_step=True, on_epoch=True, logger=True, prog_bar=True) return loss def test_step(self, test_batch, batch_idx): (x_i, _), _ = test_batch h_i, h_j, z_i, z_j = self.forward(x_i, x_i) h_i = h_i.cpu().numpy() if len(self.test_outputs) == 0: self.test_outputs = h_i else: self.test_outputs = np.append(self.test_outputs, h_i, axis=0) def test_epoch_end(self, outputs): output_csv = pd.DataFrame(self.test_outputs) output_csv.to_csv(self.args.csv_path + self.args.model_file[:-4] + "csv", header=False, index=False) # def configure_criterion(self): # criterion = NT_Xent(self.args.batch_size, self.args.temperature) # return criterion def configure_optimizers(self): scheduler = None if self.args.optimizer == "Adam": optimizer = torch.optim.Adam(self.model.parameters(), lr=self.args.learning_rate, amsgrad=self.args.amsgrad) elif self.args.optimizer == "LARS": # optimized using LARS with linear learning rate scaling # (i.e. LearningRate = 0.3 × BatchSize/256) and weight decay of 10−6. learning_rate = 0.3 * args.batch_size / 256 optimizer = LARS( self.model.parameters(), lr=learning_rate, weight_decay=args.weight_decay, exclude_from_weight_decay=["batch_normalization", "bias"], ) # "decay the learning rate with the cosine decay schedule without restarts" scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, args.epochs, eta_min=0, last_epoch=-1) else: raise NotImplementedError if scheduler: return {"optimizer": optimizer, "lr_scheduler": scheduler} else: return {"optimizer": optimizer}