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)
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 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 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 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 __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 build_model(cfg):
cfg.build_backbone = build_backbone
model = SimCLR(cfg)
logger = logging.getLogger(__name__)
logger.info("Model:\n{}".format(model))
return model
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()
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()
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)
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 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():
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():
mid = open("config.yaml", "r")
config = yaml.load(mid, Loader=yaml.FullLoader)
dataset = DataSetWrapper(config['batch_size'], **config['dataset'])
simclr = SimCLR(dataset, config)
simclr.train()
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 main():
config_file = sys.argv[1]
try:
config = load_config(config_file)
except:
print('Please provide a valid config file.')
save_parent_path = '{}_{}'.format(config_file.split(".")[0], config.title)
save_path = os.path.join(save_parent_path, config.save_path)
if not os.path.isdir(save_path): os.makedirs(save_path)
## Gather dataset
train_data = torchvision.datasets.CIFAR10(config.data_location,
train=True,
download=True,
transform=torchvision.transforms.ToTensor())
# Create dataloader
dataloader = torch.utils.data.DataLoader(train_data,
batch_size=config.batch_size,
shuffle=True,
num_workers=4)
# Define model
simclr = SimCLR(np.array(train_data[0][0]).shape[1:],
model_name=config.simclr.model_name,
embedding_size=config.simclr.embedding_size,
temp=config.simclr.temp,
transform=config.simclr.transforms,
device=device)
# Define optimizer and learning rate schedule
opt = torch.optim.Adam(simclr.parameters(),
lr=config.opt.lr,
weight_decay=config.opt.weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt,
T_max=len(dataloader))
# Start training
loss_stats = []
for epoch in tqdm.tqdm(range(config.num_epochs)):
for (x, _) in tqdm.tqdm(dataloader):
loss = simclr(x)
opt.zero_grad()
loss.backward()
opt.step()
loss_stats.append(float(loss.data.cpu().numpy()))
if epoch >=10:
scheduler.step()
print('Plotting after epoch {}/{}'.format(epoch+1, config.num_epochs))
plot_loss(loss_stats,
title=config.title + ' training curve',
path=os.path.join(save_parent_path, 'loss.png'))
# saving model after every 100 epochs
if epoch%100==0:
PATH = save_path + '{}.pth'.format(epoch)
torch.save(simclr.state_dict(), PATH)
PATH = save_path + '{}.pth'.format(epoch)
torch.save(simclr.state_dict(), PATH)
np.save(os.path.join(save_parent_path, 'loss.npy'), loss_stats, allow_pickle=True)
def main():
config = yaml.load(open('config.yaml', 'r'), Loader=yaml.FullLoader)
dataset = DataSetWrapper(config['batch_size'], **config['dataset'])
simclr = SimCLR(dataset, config)
simclr.train()
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
torch.multiprocessing.set_start_method('spawn')
else:
args.device = torch.device('cpu')
args.gpu_index = -1
if args.eval:
# Load pretrained model and cifar10
cifar_transforms = transforms.Compose(
[transforms.Resize(96),
transforms.ToTensor()])
dataset = CIFAR10(root='datasets/cifar10',
download=True,
transform=cifar_transforms)
train_dataset = CIFAR10(root='datasets/cifar10',
train=True,
transform=cifar_transforms)
valid_dataset = CIFAR10(root='datasets/cifar10',
train=False,
transform=cifar_transforms)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
valid_loader = torch.utils.data.DataLoader(valid_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
if not args.baseline:
model = ResNetBertSimCLR(base_model=args.arch,
out_dim=args.out_dim)
checkpoint = torch.load(args.saved_path)
model.load_state_dict(checkpoint['state_dict'])
else:
model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim)
checkpoint = torch.load(args.saved_path)
model.load_state_dict(checkpoint['state_dict'])
classifier_model = torch.nn.Sequential(torch.nn.Linear(128, 10))
classifier_optimizer = torch.optim.Adam(classifier_model.parameters(),
args.lr,
weight_decay=args.weight_decay)
optimizer = None
scheduler = None
else:
# Load BertSimCLR and coco dataset
dataset = ContrastiveLearningDataset(args.data)
train_dataset = dataset.get_dataset(args.dataset_name, args.n_views)
valid_dataset = dataset.get_dataset(args.dataset_name + 'valid',
args.n_views)
if not args.baseline:
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,
collate_fn=coco_collate_fn)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
drop_last=True,
collate_fn=coco_collate_fn)
model = ResNetBertSimCLR(base_model=args.arch,
out_dim=args.out_dim)
else:
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)
valid_loader = torch.utils.data.DataLoader(
valid_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)
# emergency fix
"""
model.to(args.device)
checkpoint = torch.load('/home/gabriel/Desktop/Projects/SimCLR/runs/nextbase/checkpoint_0018.pth.tar', map_location="cuda:0")
model_state = checkpoint['state_dict']#.to(args.device)
opt_state = checkpoint['optimizer']#.to(args.device)
model.load_state_dict(model_state)
optimizer.load_state_dict(opt_state)
model.to(args.device)
"""
classifier_model = None
classifier_optimizer = None
#print(model.visual_backbone)
data_loaders = {"train": train_loader, "val": valid_loader}
# It’s a no-op if the 'gpu_index' argument is a negative integer or None.
with torch.cuda.device(args.gpu_index):
if not args.baseline:
simclr = BertSimCLR(model=model,
optimizer=optimizer,
scheduler=scheduler,
classifier_model=classifier_model,
classifier_optimizer=classifier_optimizer,
args=args)
if args.eval:
simclr.train_linear_classifier(args.epochs, data_loaders)
else:
simclr.train(data_loaders)
else:
simclr = SimCLR(model=model,
optimizer=optimizer,
scheduler=scheduler,
classifier_model=classifier_model,
classifier_optimizer=classifier_optimizer,
args=args)
if args.eval:
simclr.train_linear_classifier(args.epochs, data_loaders)
else:
simclr.train(data_loaders)
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(args, encoder, 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()
def main(gpu, args):
rank = args.nr * args.gpus + gpu
if args.nodes > 1:
dist.init_process_group("nccl", rank=rank, world_size=args.world_size)
torch.cuda.set_device(gpu)
torch.manual_seed(args.seed)
np.random.seed(args.seed)
if args.dataset == "STL10":
train_dataset = torchvision.datasets.STL10(
args.dataset_dir,
split="unlabeled",
download=True,
transform=TransformsSimCLR(size=args.image_size),
)
elif args.dataset == "CIFAR10":
train_dataset = torchvision.datasets.CIFAR10(
args.dataset_dir,
download=True,
transform=TransformsSimCLR(size=args.image_size),
)
else:
raise NotImplementedError
if args.nodes > 1:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset,
num_replicas=args.world_size,
rank=rank,
shuffle=True)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
drop_last=True,
num_workers=args.workers,
sampler=train_sampler,
)
# initialize ResNet
encoder = get_resnet(args.resnet, pretrained=False)
n_features = encoder.fc.in_features # get dimensions of fc layer
# initialize model
model = SimCLR(encoder, args.projection_dim, n_features)
if args.reload:
model_fp = os.path.join(args.model_path,
"checkpoint_{}.tar".format(args.epoch_num))
model.load_state_dict(
torch.load(model_fp, map_location=args.device.type))
model = model.to(args.device)
# optimizer / loss
optimizer, scheduler = load_optimizer(args, model)
criterion = NT_Xent(args.batch_size, args.temperature, args.world_size)
# DDP / DP
if args.dataparallel:
model = convert_model(model)
model = DataParallel(model)
else:
if args.nodes > 1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = DDP(model, device_ids=[gpu])
model = model.to(args.device)
writer = None
if args.nr == 0:
writer = SummaryWriter()
args.global_step = 0
args.current_epoch = 0
for epoch in range(args.start_epoch, args.epochs):
lr = optimizer.param_groups[0]["lr"]
loss_epoch = train(args, train_loader, model, criterion, optimizer,
writer)
if args.nr == 0 and scheduler:
scheduler.step()
if args.nr == 0 and epoch % 10 == 0:
save_model(args, model, optimizer)
if args.nr == 0:
writer.add_scalar("Loss/train", loss_epoch / len(train_loader),
epoch)
writer.add_scalar("Misc/learning_rate", lr, epoch)
print(
f"Epoch [{epoch}/{args.epochs}]\t Loss: {loss_epoch / len(train_loader)}\t lr: {round(lr, 5)}"
)
args.current_epoch += 1
## end training
save_model(args, model, optimizer)
