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}
