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 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 predict(self):
#Data
valid_loader = self.dataset.get_data_loaders()
#Model
model = ResNetSimCLR(**self.config["model"])
if self.device == 'cuda':
model = nn.DataParallel(
model,
device_ids=[i for i in range(self.config['gpu']['gpunum'])])
model = model.cuda()
#print(model)
model = self._load_pre_trained_weights(model)
# validate the model if requested
self._validate(model, valid_loader)
def __init__(self, dataset, config):
self.config = config
self.device = self._get_device()
self.writer = SummaryWriter()
self.dataset = dataset
self.model = ResNetSimCLR(**self.config["model"]).to(self.device)
self.criterion = self._make_loss()
def __init__(self, dataset, config):
self.config = config
self.device = self._get_device()
self.writer = SummaryWriter()
self.dataset = dataset
self.batch_size = config['batch_size']
self.model = ResNetSimCLR(**self.config["model"]).to(self.device)
self.loss_func = self._choose_loss()
def load_model(checkpoints_folder,device):
model =ResNetSimCLR(**config['model'])
model.eval()
state_dict = torch.load(os.path.join(checkpoints_folder, 'model.pth'), map_location=torch.device('cpu'))
model.load_state_dict(state_dict)
model = model.to(device)
return model
def _build_model(self):
model = ResNetSimCLR(**self.config["model"]).to(self.device)
model = self._load_pre_trained_weights(model)
augmentor = IcmAugmentorv2(num_mech=self.config["num_mechanisms"]).to(
self.device)
augmentor = self._load_pre_trained_weights(augmentor, "augmentor.pth")
discriminator = Discriminatorv2(
num_mech=self.config["num_mechanisms"]).to(self.device)
discriminator = self._load_pre_trained_weights(discriminator,
"discriminator.pth")
return model, augmentor, discriminator
def encode(save_root, model_file, data_folder, model_name='ca', dataset_name='celeba', batch_size=64, device='cuda:0', out_dim=256):
os.makedirs(save_root, exist_ok=True)
os.makedirs(data_folder, exist_ok=True)
if dataset_name == 'celeba':
train_loader = DataLoader(datasets.CelebA(data_folder, split='train', download=True, transform=transforms.ToTensor()),
batch_size=batch_size, shuffle=False)
valid_loader = DataLoader(datasets.CelebA(data_folder, split='valid', download=True, transform=transforms.ToTensor()),
batch_size=batch_size, shuffle=False)
elif dataset_name == 'stanfordCars':
t = transforms.Compose([
transforms.Resize(512),
transforms.CenterCrop(512),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.repeat(3,1,1) if x.shape[0] == 1 else x)
])
train_data_dir = os.path.join(data_folder, 'cars_train/')
train_annos = os.path.join(data_folder, 'devkit/cars_train_annos.mat')
train_loader = DataLoader(CarsDataset(train_annos, train_data_dir, t), batch_size=batch_size, shuffle=False)
valid_data_dir = os.path.join(data_folder, 'cars_test/')
valid_annos = os.path.join(data_folder, 'devkit/cars_test_annos_withlabels.mat')
valid_loader = DataLoader(CarsDataset(valid_annos, valid_data_dir, t), batch_size=batch_size, shuffle=False)
elif dataset_name == 'compCars':
t = transforms.Compose([
transforms.Resize(512),
transforms.CenterCrop(512),
transforms.ToTensor()
])
train_loader = DataLoader(CompCars(data_folder, True, t), batch_size=batch_size, shuffle=False)
valid_loader = DataLoader(CompCars(data_folder, False, t), batch_size=batch_size, shuffle=False)
model = ResNetSimCLR('resnet50', out_dim)
model.load_state_dict(torch.load(model_file, map_location=device))
model = model.to(device)
model.eval()
print('Starting on training data')
train_encodings = []
for x, _ in train_loader:
x = x.to(device)
h, _ = model(x)
train_encodings.append(h.cpu().detach())
torch.save(torch.cat(train_encodings, dim=0), os.path.join(save_root, f'{dataset_name}-{model_name}model-train_encodings.pt'))
print('Starting on validation data')
valid_encodings = []
for x, _ in valid_loader:
x = x.to(device)
h, _ = model(x)
if len(h.shape) == 1:
h = h.unsqueeze(0)
valid_encodings.append(h.cpu().detach())
torch.save(torch.cat(valid_encodings, dim=0), os.path.join(save_root, f'{dataset_name}-{model_name}model-valid_encodings.pt'))
def _save_dt_features(self):
from eval.feature_saver import LvisSaver
import torch
config = self.config
if self.config['hyperbolic']:
from models.hyperbolic_resnet import HResNetSimCLR
model = HResNetSimCLR(config['model']['base_model'], config['model']['out_dim'])
else:
from models.resnet_simclr import ResNetSimCLR
model = ResNetSimCLR(config['model']['base_model'], config['model']['out_dim'])
state_dict = torch.load(os.path.join(self.run_path, self.model_ckpt)) # , map_location=device)i
model.load_state_dict(state_dict)
model.eval()
saver = LvisSaver(model, self.lvis_dt, GT_FEATS)
saver.save()
def main():
config_path = "config.json"
args = process_config(config_path)
# 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
log_dir = time.ctime()
os.makedirs(log_dir)
model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim)
load_checkpoint(model, args.simclr.pretrained_path)
# run_simclr(args, model, log_dir)
run_la(args, model, log_dir)
def _init_model_and_optimizer(self):
if self.config['hyperbolic']:
model = HResNetSimCLR(**self.config["model"]).to(self.device)
else:
model = ResNetSimCLR(**self.config["model"]).to(self.device)
model, loaded_iter = self._load_pre_trained_weights(model)
# if self.config['hyperbolic']:
# optimizer = geoopt.optim.RiemannianAdam(
# [p for p in self.model.parameters() if p.requires_grad],
# 1e-4, weight_decay=eval(self.config['weight_decay']))
# else:
lr = float(self.config['lr'])
optimizer = torch.optim.Adam(
[p for p in model.parameters() if p.requires_grad],
lr,
weight_decay=eval(self.config['weight_decay']))
self.model = model
self.optimizer = optimizer
return loaded_iter
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
set_random_seed(args.seed)
train_loader, valid_loader = get_dataloaders(args)
if args.mode == 'simclr':
model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim)
trainer_class = SimCLRTrainer
elif args.mode == 'supervised':
model = ResNetSimCLR(base_model=args.arch, out_dim=len(train_loader.dataset.classes))
trainer_class = SupervisedTrainer
else:
raise InvalidTrainingMode()
if args.optimizer_mode == 'simclr':
optimizer = torch.optim.Adam(model.parameters(), lr=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)
else:
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
momentum=0.9, weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
# It’s a no-op if the 'gpu_index' argument is a negative integer or None.
with torch.cuda.device(args.gpu_index):
trainer = trainer_class(model=model, optimizer=optimizer, scheduler=scheduler, args=args)
trainer.train(train_loader, valid_loader)
def run(config):
model = ResNetSimCLR('resnet50', config.out_dim)
model.load_state_dict(
torch.load(config.model_file, map_location=config.device))
model = model.to(config.device)
clf = nn.Linear(2048, 196)
full = FineTuner(model, clf)
optim = torch.optim.Adam(list(model.parameters()) + list(clf.parameters()))
objective = nn.CrossEntropyLoss()
t = T.Compose([
T.Resize(512),
T.CenterCrop(512),
T.ToTensor(),
T.Lambda(lambda x: x.repeat(3, 1, 1) if x.shape[0] == 1 else x)
])
train_data_dir = os.path.join(config.data_root, 'cars_train/')
train_annos = os.path.join(config.data_root, 'devkit/cars_train_annos.mat')
train_dataset = StanfordCarsMini(train_annos, train_data_dir, t)
train_loader = DataLoader(train_dataset,
batch_size=config.batch_size,
shuffle=True)
valid_data_dir = os.path.join(config.data_root, 'cars_test/')
valid_annos = os.path.join(config.data_root,
'devkit/cars_test_annos_withlabels.mat')
valid_dataset = CarsDataset(valid_annos, valid_data_dir, t)
valid_loader = DataLoader(valid_dataset, batch_size=config.batch_size)
solver = CESolver(full,
train_loader,
valid_loader,
config.save_root,
name=config.name,
device=config.device)
solver.train(config.num_epochs)
def eval_frozen(self):
train_loader, val_loader, num_classes = self.dataset.get_dataset_eval()
model = ResNetSimCLR(**self.config["model"]).to(self.device)
model = self._load_pre_trained_weights(model)
model.to(self.device)
model.eval()
lineal_classifier = nn.Linear(model.l1.in_features, num_classes)
lineal_classifier.to(self.device)
optimizer = torch.optim.SGD(lineal_classifier.parameters(),
1e-3,
weight_decay=eval(
self.config['weight_decay']))
epochs = self.config['epochs']
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[int(0.5 * epochs),
int(0.8 * epochs)],
gamma=0.1,
last_epoch=-1)
criterion = nn.CrossEntropyLoss()
best_acc = 0
for epoch in range(epochs):
print('=' * 20)
print(f'Epoch: {epoch + 1} / {epochs}')
top1 = 0
running_loss = 0
n = 0
lineal_classifier.train()
for idx, (img, lab) in enumerate(train_loader):
B = img.size(0)
img = img.to(self.device)
lab = lab.to(self.device)
optimizer.zero_grad()
loss, top1_batch = self._step_eval_train(
model, lineal_classifier, img, lab)
loss.backward()
optimizer.step()
top1 += (x.argmax(dim=1) == lab).sum().item()
running_loss += loss.item() * B
n += B
print('Training {}/{} - Loss: {:.2f} - top 1: {:.2f}'.format(
idx + 1, len(train_loader), running_loss / n,
100 * top1 / n),
end='\r')
print('\n')
top1 = 0
running_loss = 0
n = 0
lineal_classifier.eval()
for idx, (img, lab) in enumerate(val_loader):
B = img.size(0)
img = img.to(self.device)
lab = lab.to(self.device)
loss, top1_batch = self._step_eval_eval(
model, lineal_classifier, img, lab)
top1 += top1_batch
running_loss += loss.item() * B
n += B
print('Val {}/{} - Loss: {:.2f} - top 1: {:.2f}'.format(
idx + 1, len(val_loader), running_loss / n,
100 * top1 / n),
end='\r')
print('\n')
if best_acc < top1:
best_acc = top1
print(f'Best ACC: {best_acc * 100 / n}')
def train(self):
train_loader, valid_loader = self.dataset.get_data_loaders()
model = ResNetSimCLR(**self.config["model"]).to(self.device)
model = self._load_pre_trained_weights(model)
optimizer = torch.optim.Adam(model.parameters(),
3e-4,
weight_decay=eval(
self.config['weight_decay']))
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=self.config['epochs'], eta_min=0, last_epoch=-1)
if apex_support and self.config['fp16_precision']:
model, optimizer = amp.initialize(model,
optimizer,
opt_level='O2',
keep_batchnorm_fp32=True)
model_checkpoints_folder = os.path.join(self.writer.log_dir,
'checkpoints')
# save config file
_save_config_file(model_checkpoints_folder)
n_iter = 0
valid_n_iter = 0
best_valid_loss = np.inf
for epoch_counter in range(self.config['epochs']):
for (xis, xjs), _ in train_loader:
optimizer.zero_grad()
xis = xis.to(self.device)
xjs = xjs.to(self.device)
loss = self._step(model, xis, xjs, n_iter)
if n_iter % self.config['log_every_n_steps'] == 0:
self.writer.add_scalar('train_loss',
loss,
global_step=n_iter)
if apex_support and self.config['fp16_precision']:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
n_iter += 1
# validate the model if requested
if epoch_counter % self.config['eval_every_n_epochs'] == 0:
valid_loss = self._validate(model, valid_loader)
if valid_loss < best_valid_loss:
# save the model weights
best_valid_loss = valid_loss
torch.save(
model.state_dict(),
os.path.join(model_checkpoints_folder, 'model.pth'))
self.writer.add_scalar('validation_loss',
valid_loss,
global_step=valid_n_iter)
valid_n_iter += 1
# warmup for the first 10 epochs
if epoch_counter >= 10:
scheduler.step()
self.writer.add_scalar('cosine_lr_decay',
scheduler.get_lr()[0],
global_step=n_iter)
def run(config):
model = ResNetSimCLR('resnet50', config.out_dim)
model.load_state_dict(
torch.load(config.model_file, map_location=config.device))
model = model.to(config.device)
clf = nn.Linear(2048, 196)
train_data_dir = os.path.join(config.data_root, 'cars_train/')
train_annos = os.path.join(config.data_root, 'devkit/cars_train_annos.mat')
valid_data_dir = os.path.join(config.data_root, 'cars_test/')
valid_annos = os.path.join(config.data_root,
'devkit/cars_test_annos_withlabels.mat')
if config.encodings_file_prefix:
train_dataset = EncodedStanfordCarsDataset(
train_annos, config.encodings_file_prefix + '-train_encodings.pt')
train_loader = DataLoader(train_dataset,
batch_size=config.encodings_batch_size,
shuffle=True)
valid_dataset = EncodedStanfordCarsDataset(
valid_annos, config.encodings_file_prefix + '-valid_encodings.pt')
valid_loader = DataLoader(valid_dataset,
batch_size=config.encodings_batch_size)
tmp_clf = nn.Linear(2048, 196)
clf_solver = CESolver(tmp_clf,
train_loader,
valid_loader,
config.save_root,
name=config.name + '-clf',
device=config.device)
clf_solver.train(config.encodings_num_epochs)
clf_filename = os.path.join(config.save_root, f'{config.name}-clf.pth')
clf.load_state_dict(
torch.load(clf_filename, map_location=config.device))
full = FineTuner(model, clf)
t = T.Compose([
T.Resize(512),
T.CenterCrop(512),
T.ToTensor(),
T.Lambda(lambda x: x.repeat(3, 1, 1) if x.shape[0] == 1 else x)
])
train_dataset = StanfordCarsMini(train_annos, train_data_dir, t)
train_loader = DataLoader(train_dataset,
batch_size=config.batch_size,
shuffle=True)
valid_dataset = CarsDataset(valid_annos, valid_data_dir, t)
valid_loader = DataLoader(valid_dataset, batch_size=config.batch_size)
solver = CESolver(full,
train_loader,
valid_loader,
config.save_root,
name=config.name,
device=config.device)
# full = full.to(config.device)
# print(solver.validate(full, nn.CrossEntropyLoss()))
solver.train(config.num_epochs)
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
if args.mode == 'simclr':
dataset = ContrastiveLearningDataset(args.data)
train_dataset = dataset.get_dataset(args.dataset_name,
args.n_views,
train=True)
model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim)
trainer_class = SimCLRTrainer
elif args.mode == 'supervised':
dataset = SupervisedLearningDataset(args.data)
train_dataset = dataset.get_dataset(args.dataset_name,
args.supervised_augments,
train=True)
model = ResNetSimCLR(base_model=args.arch,
out_dim=len(train_dataset.classes))
trainer_class = SupervisedTrainer
else:
raise InvalidTrainingMode()
if args.target_shuffle is not None:
random.seed(args.target_shuffle)
random.shuffle(train_dataset.targets)
checkpoints = []
for root, dirs, files in os.walk(
os.path.join('experiments', args.experiment_group, 'wandb')):
for file in files:
if file == args.estimate_checkpoint:
checkpoints += [os.path.join(root, file)]
set_random_seed(args.seed)
sample_indices = torch.randint(len(train_dataset),
size=(args.batch_size *
args.estimate_batches, ))
# It’s a no-op if the 'gpu_index' argument is a negative integer or None.
estimated_prob, estimated_argmax = [], []
with torch.cuda.device(args.gpu_index):
for file in checkpoints:
state = torch.load(file)
model.load_state_dict(state['model'])
model.eval()
trainer = trainer_class(model=model,
optimizer=None,
scheduler=None,
args=args)
checkpoint_prob, checkpoint_argmax = [], []
for i in range(args.estimate_batches):
if args.fixed_augments:
set_random_seed(args.seed)
if args.mode == 'simclr':
images = [[], []]
for index in sample_indices[i:i + args.batch_size]:
example = train_dataset[index][0]
images[0] += [example[0]]
images[1] += [example[1]]
images[0] = torch.stack(images[0], dim=0)
images[1] = torch.stack(images[1], dim=0)
labels = None
elif args.mode == 'supervised':
images, labels = [], []
for index in sample_indices[i:i + args.batch_size]:
example = train_dataset[index]
images += [example[0]]
labels += [example[1]]
images = torch.stack(images, dim=0)
labels = torch.tensor(labels, dtype=torch.long)
with torch.no_grad():
logits, labels = trainer.calculate_logits(images, labels)
prob = torch.softmax(logits,
dim=1)[torch.arange(labels.shape[0]),
labels]
checkpoint_prob += [prob.detach().cpu()]
argmax = (torch.argmax(logits,
dim=1) == labels).to(torch.int)
checkpoint_argmax += [argmax.detach().cpu()]
checkpoint_prob = torch.cat(checkpoint_prob, dim=0)
estimated_prob += [checkpoint_prob]
checkpoint_argmax = torch.cat(checkpoint_argmax, dim=0)
estimated_argmax += [checkpoint_argmax]
estimated_prob = torch.stack(estimated_prob, dim=0)
estimated_argmax = torch.stack(estimated_argmax, dim=0)
torch.save(
{
'indices': sample_indices,
'prob': estimated_prob,
'argmax': estimated_argmax
}, os.path.join('experiments', args.experiment_group, args.out_file))
test_loss = (test_loss * i + criterion(out, labels).item()) / (i + 1)
res = torch.argmax(out, 1)
correct += (res == labels).sum().item()
tot += labels.size(0)
return 100 - 100 * correct / tot, test_loss
config_path = "./best_run/config.yaml"
model_path = "./best_run/model.pth"
data_path = "./data"
config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader)
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
model = ResNetSimCLR(**config["model"]).to(device)
model.load_state_dict(torch.load(model_path))
lr = 0.001
num_epoch = 90
batch_size = 512
num_classes = 10
weight_decay = 1e-6
transform = transforms.Compose([
transforms.ToTensor(),
])
linear = Linear_Model(512, num_classes) # Linear_Model or MLP
linear.to(device)
def train(self):
#Data
train_loader, valid_loader = self.dataset.get_data_loaders()
#Model
model = ResNetSimCLR(**self.config["model"])
if self.device == 'cuda':
model = nn.DataParallel(model, device_ids=[i for i in range(self.config['gpu']['gpunum'])])
#model = model.to(self.device)
model = model.cuda()
print(model)
model = self._load_pre_trained_weights(model)
each_epoch_steps = len(train_loader)
total_steps = each_epoch_steps * self.config['train']['epochs']
warmup_steps = each_epoch_steps * self.config['train']['warmup_epochs']
scaled_lr = eval(self.config['train']['lr']) * self.batch_size / 256.
optimizer = torch.optim.Adam(
model.parameters(),
scaled_lr,
weight_decay=eval(self.config['train']['weight_decay']))
'''
optimizer = LARS(params=model.parameters(),
lr=eval(self.config['train']['lr']),
momentum=self.config['train']['momentum'],
weight_decay=eval(self.config['train']['weight_decay'],
eta=0.001,
max_epoch=self.config['train']['epochs'])
'''
# scheduler during warmup stage
lambda1 = lambda epoch:epoch*1.0 / int(warmup_steps)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1)
if apex_support and self.config['train']['fp16_precision']:
model, optimizer = amp.initialize(model, optimizer,
opt_level='O2',
keep_batchnorm_fp32=True)
model_checkpoints_folder = os.path.join(self.writer.log_dir, 'checkpoints')
# save config file
_save_config_file(model_checkpoints_folder)
n_iter = 0
valid_n_iter = 0
best_valid_loss = np.inf
lr = eval(self.config['train']['lr'])
end = time.time()
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
for epoch_counter in range(self.config['train']['epochs']):
model.train()
for i, ((xis, xjs), _) in enumerate(train_loader):
data_time.update(time.time() - end)
optimizer.zero_grad()
xis = xis.cuda()
xjs = xjs.cuda()
loss = self._step(model, xis, xjs, n_iter)
#print("Loss: ",loss.data.cpu())
losses.update(loss.item(), 2 * xis.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
print('Epoch: [{epoch}][{step}/{each_epoch_steps}] Loss {loss.val:.4f} Avg Loss {loss.avg:.4f} DataTime {datatime.val:.4f} BatchTime {batchtime.val:.4f} LR {lr})'.format(epoch=epoch_counter, step=i, each_epoch_steps=each_epoch_steps, loss=losses, datatime=data_time, batchtime=batch_time, lr=lr))
if n_iter % self.config['train']['log_every_n_steps'] == 0:
self.writer.add_scalar('train_loss', loss, global_step=n_iter)
if apex_support and self.config['train']['fp16_precision']:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
n_iter += 1
#adjust lr
if n_iter == warmup_steps:
# scheduler after warmup stage
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps-warmup_steps, eta_min=0, last_epoch=-1)
scheduler.step()
lr = scheduler.get_lr()[0]
self.writer.add_scalar('cosine_lr_decay', scheduler.get_lr()[0], global_step=n_iter)
sys.stdout.flush()
# validate the model if requested
if epoch_counter % self.config['train']['eval_every_n_epochs'] == 0:
valid_loss = self._validate(model, valid_loader)
if valid_loss < best_valid_loss:
# save the model weights
best_valid_loss = valid_loss
torch.save(model.state_dict(), os.path.join(model_checkpoints_folder, 'model.pth'))
self.writer.add_scalar('validation_loss', valid_loss, global_step=valid_n_iter)
valid_n_iter += 1
# dataset_train, batch_size=args.batch_size, shuffle=True,
# num_workers=args.workers, pin_memory=True, drop_last=True)
# build model
if args.model == 'cnn' and args.dataset == 'cifar':
net_glob = CNNCifar(args=args).to(args.device)
elif args.model == 'cnn' and args.dataset == 'mnist':
net_glob = CNNMnist(args=args).to(args.device)
elif args.model == 'mlp':
len_in = 1
for x in img_size:
len_in *= x
net_glob = MLP(dim_in=len_in, dim_hidden=200,
dim_out=args.num_classes).to(args.device)
elif args.model == 'resnet':
net_glob = ResNetSimCLR(base_model=args.arch,
out_dim=args.out_dim).to(args.device)
else:
exit('Error: unrecognized model')
net_glob.train()
# copy weights
w_glob = net_glob.state_dict()
dict_glob = {}
# training
loss_train = []
cv_loss, cv_acc = [], []
val_loss_pre, counter = 0, 0
net_best = None
best_loss = None
def main_worker(gpu, ngpus_per_node, args):
args.gpu = gpu
# suppress printing if not master
if args.multiprocessing_distributed and args.gpu != 0:
def print_pass(*args):
pass
builtins.print = print_pass
if args.gpu is not None:
print(args.gpu)
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create model
print("=> creating model '{}'".format(args.arch))
model = ResNetSimCLR(base_model=args.arch,
out_dim=args.out_dim).to(args.gpu)
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.num_workers = int(
(args.num_workers + ngpus_per_node - 1) / ngpus_per_node)
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
# comment out the following line for debugging
#raise NotImplementedError("Only DistributedDataParallel is supported.")
#else:
# AllGather implementation (batch shuffle, queue update, etc.) in
# this code only supports DistributedDataParallel.
#raise NotImplementedError("Only DistributedDataParallel is supported.")
# Data loader
train_loader, train_sampler = data_loader(args.dataset,
args.data_path,
args.batch_size,
args.num_workers,
download=args.download,
distributed=args.distributed,
supervised=False)
#optimizer = torch.optim.Adam(model.parameters(), 3e-4, weight_decay=args.weight_decay)
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=args.epochs,
eta_min=0,
last_epoch=-1)
criterion = NTXentLoss(args.gpu, args.batch_size, args.temperature,
True).cuda(args.gpu)
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
if args.gpu is None:
checkpoint = torch.load(args.resume)
else:
loc = 'cuda:{}'.format(args.gpu)
checkpoint = torch.load(args.resume, map_location=loc)
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
if apex_support and args.fp16_precision:
model, optimizer = amp.initialize(model,
optimizer,
opt_level='O2',
keep_batchnorm_fp32=True)
cudnn.benchmark = True
train(model, train_loader, train_sampler, criterion, optimizer, scheduler,
args, ngpus_per_node)
train_dataset = datasets.STL10('./data',
split='train+unlabeled',
download=True,
transform=DataTransform(data_augment))
# train_dataset = datasets.Caltech101(root='./data', target_type="category", transform=transforms.ToTensor(),
# target_transform=None, download=True)
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
num_workers=config['num_workers'],
drop_last=True,
shuffle=True)
# model = Encoder(out_dim=out_dim)
model = ResNetSimCLR(base_model=config["base_convnet"], out_dim=out_dim)
train_gpu = torch.cuda.is_available()
print("Is gpu available:", train_gpu)
# moves the model parameters to gpu
if train_gpu:
model.cuda()
criterion = torch.nn.CrossEntropyLoss(reduction='sum')
optimizer = optim.Adam(model.parameters(), 3e-4)
train_writer = SummaryWriter()
sim_func_dim1, sim_func_dim2 = get_similarity_function(use_cosine_similarity)
def train(self):
train_loader, valid_loader = self.dataset.get_data_loaders()
model = ResNetSimCLR(**self.config["model"]).to(self.device) #just a resnet backbone
model = self._load_pre_trained_weights(model) #checkpoints (shall we convert TF checkpoint in to Torch and train or train from the scratch since we have f*****g lot?)
optimizer = torch.optim.Adam(model.parameters(), 3e-4, weight_decay=eval(self.config['weight_decay']))
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0,
last_epoch=-1) #learning rate shedulers (let's use as it is)
if apex_support and self.config['fp16_precision']:
model, optimizer = amp.initialize(model, optimizer,
opt_level='O2',
keep_batchnorm_fp32=True)
model_checkpoints_folder = os.path.join(self.writer.log_dir, 'checkpoints')
# save config file
_save_config_file(model_checkpoints_folder)
n_iter = 0
valid_n_iter = 0
best_valid_loss = np.inf
for epoch_counter in range(self.config['epochs']): #start training
for (x, y) in train_loader: #dataset
optimizer.zero_grad()
x = x.to(self.device) #in SimCLR we calculate the loss with two augmentation versions
y = y.to(self.device)
loss = self._step(model, x, y)
if n_iter % self.config['log_every_n_steps'] == 0:
self.writer.add_scalar('train_loss', loss, global_step=n_iter)
if apex_support and self.config['fp16_precision']:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
n_iter += 1
# validate the model if requested
if epoch_counter % self.config['eval_every_n_epochs'] == 0:
valid_loss = self._validate(model, valid_loader)
print('Epoch:',epoch_counter,' ---',' validation_loss:',valid_loss)
if valid_loss < best_valid_loss:
# save the model weights
best_valid_loss = valid_loss
torch.save(model.state_dict(), os.path.join(model_checkpoints_folder, 'model.pth'))
self.writer.add_scalar('validation_loss', valid_loss, global_step=valid_n_iter)
valid_n_iter += 1
# warmup for the first 10 epochs
if epoch_counter >= 10:
scheduler.step()
self.writer.add_scalar('cosine_lr_decay', scheduler.get_lr()[0], global_step=n_iter)
def train(self):
train_loader, valid_loader = self.dataset.get_data_loaders()
model = ResNetSimCLR(**self.config["model"]).to(self.device)
model = self._load_pre_trained_weights(model)
criterion = nn.CrossEntropyLoss() # loss function
optimizer = torch.optim.Adam(model.parameters(),
3e-4,
weight_decay=eval(
self.config['weight_decay']))
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=self.config['epochs'], eta_min=0, last_epoch=-1)
if apex_support and self.config['fp16_precision']:
model, optimizer = amp.initialize(model,
optimizer,
opt_level='O2',
keep_batchnorm_fp32=True)
model_checkpoints_folder = os.path.join(
'/home/zhangchunhui/MedicalAI/USCL/checkpoints_multi_aug',
'checkpoint_' + str(self.Checkpoint_Num))
# save config file
_save_config_file(model_checkpoints_folder)
start_time = time.time()
end_time = time.time()
valid_n_iter = 0
best_valid_loss = np.inf
for epoch in range(self.config['epochs']):
for i, data in enumerate(train_loader, 1):
# forward
# mixupimg1, label1, mixupimg2, label2, original img1, original img2
xis, labelis, xjs, labeljs, imgis, imgjs = data # N samples of left branch, N samples of right branch
xis = xis.to(self.device)
xjs = xjs.to(self.device)
####### 1-Semi-supervised
hi, xi, outputis = model(xis)
hj, xj, outputjs = model(xjs)
labelindexi, labelindexj = FindNotX(
labelis.tolist(), 9999), FindNotX(labeljs.tolist(),
9999) # X=9999=no label
lossi = criterion(outputis[labelindexi],
labelis.to(self.device)[labelindexi])
lossj = criterion(outputjs[labelindexj],
labeljs.to(self.device)[labelindexj])
# lumbda1=lumbda2 # small value is better
lumbda1, lumbda2 = self.lumbda1, self.lumbda2 # small value is better
loss = self._step(model, xis,
xjs) + lumbda1 * lossi + lumbda2 * lossj
########################################################################################################
####### 2-Self-supervised
# loss = self._step(model, xis, xjs)
########################################################################################################
# backward
optimizer.zero_grad()
if apex_support and self.config['fp16_precision']:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
# update weights
optimizer.step()
if i % self.config['log_every_n_steps'] == 0:
# self.writer.add_scalar('train_loss', loss, global_step=i)
start_time, end_time = end_time, time.time()
print(
"\nTraining:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Time: {:.2f}s"
.format(epoch + 1, self.config['epochs'], i,
len(train_loader), loss,
end_time - start_time))
# validate the model if requested
if epoch % self.config['eval_every_n_epochs'] == 0:
start_time = time.time()
valid_loss = self._validate(model, valid_loader)
end_time = time.time()
if valid_loss < best_valid_loss:
# save the model weights
best_valid_loss = valid_loss
torch.save(
model.state_dict(),
os.path.join(model_checkpoints_folder,
'best_model.pth'))
print(
"Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Time: {:.2f}s"
.format(epoch + 1, self.config['epochs'],
len(valid_loader), len(valid_loader), valid_loss,
end_time - start_time))
# self.writer.add_scalar('validation_loss', valid_loss, global_step=valid_n_iter)
valid_n_iter += 1
print('Learning rate this epoch:',
scheduler.get_last_lr()[0]) # python >=3.7
# print('Learning rate this epoch:', scheduler.base_lrs[0]) # python 3.6
# warmup for the first 10 epochs
if epoch >= 10:
scheduler.step()
out_dim = config['out_dim']
temperature = config['temperature']
use_cosine_similarity = config['use_cosine_similarity']
data_augment = get_simclr_data_transform(s=config['s'], crop_size=96)
train_dataset = datasets.STL10('./data',
split='train+unlabeled',
download=True,
transform=DataTransform(data_augment))
train_loader, valid_loader = get_train_validation_data_loaders(
train_dataset, **config)
# model = Encoder(out_dim=out_dim)
model = ResNetSimCLR(base_model=config["base_convnet"], out_dim=out_dim)
if config['continue_training']:
checkpoints_folder = os.path.join('./runs', config['continue_training'],
'checkpoints')
state_dict = torch.load(os.path.join(checkpoints_folder, 'model.pth'))
model.load_state_dict(state_dict)
print("Loaded pre-trained model with success.")
train_gpu = torch.cuda.is_available()
print("Is gpu available:", train_gpu)
# moves the model parameters to gpu
if train_gpu:
model = model.cuda()
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)
def evaluation(checkpoints_folder, config, device):
model = ResNetSimCLR(**config['model'])
model.eval()
model.load_state_dict(
torch.load(os.path.join(checkpoints_folder, 'model.pth')))
model = model.to(device)
train_set = torchvision.datasets.CIFAR10(
root='../data/CIFAR10',
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]),
train=True,
download=True)
test_set = torchvision.datasets.CIFAR10(root='../data/CIFAR10',
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]),
train=False,
download=True)
# num_train = len(train_dataset)
# indices = list(range(num_train))
# np.random.shuffle(indices)
#
# split = int(np.floor(0.05 * num_train))
# train_idx, test_idx = indices[split:], indices[:split]
#
# # define samplers for obtaining training and validation batches
# train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_idx)
# test_sampler = torch.utils.data.sampler.SubsetRandomSampler(test_idx) # ?????sampler????????????
# train_loader = torch.utils.data.DataLoader(train_set, batch_size=config['batch_size'], drop_last=True, shuffle=True)
#
# test_loader = torch.utils.data.DataLoader(test_set, batch_size=config['batch_size'], drop_last=True, shuffle=True)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=48,
drop_last=True,
shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=48,
drop_last=True,
shuffle=True)
X_train_feature = []
label_train = []
for data in train_loader:
x, y = data
x = x.to(device)
features, _ = model(x)
X_train_feature.extend(features.cpu().detach().numpy())
label_train.extend(y.cpu().detach().numpy())
X_train_feature = np.array(X_train_feature)
label_train = np.array(label_train)
X_test_feature = []
label_test = []
for data in test_loader:
x, y = data
x = x.to(device)
features, _ = model(x)
X_test_feature.extend(features.cpu().detach().numpy())
label_test.extend(y.cpu().detach().numpy())
X_test_feature = np.array(X_test_feature)
label_test = np.array(label_test)
scaler = preprocessing.StandardScaler()
print('ok')
scaler.fit(X_train_feature)
# print(X_test_feature.shape)
# print(y_test.shape)
linear_model_eval(scaler.transform(X_train_feature), label_train,
scaler.transform(X_test_feature), label_test)
def train(self):
train_loader, valid_loader = self.dataset.get_data_loaders()
model = ResNetSimCLR(**self.config["model"]).to(self.device)
model = self._load_pre_trained_weights(model)
if self.augmentor_type == "cnn":
if self.config["normalization_type"] == "original":
augmentor = LpAugmentor(
clip=self.config["augmentor_clip_output"])
augmentor.to(self.device)
elif self.config["normalization_type"] == "spectral":
augmentor = LpAugmentorSpecNorm(
clip=self.config["augmentor_clip_output"])
augmentor.to(self.device)
else:
raise ValueError("Unregonized normalization type: {}".format(
self.config["normalization_type"]))
elif self.augmentor_type == "style_transfer":
augmentor = LpAugmentorStyleTransfer(
clip=self.config["augmentor_clip_output"])
augmentor.to(self.device)
elif self.augmentor_type == "transformer":
augmentor = LpAugmentorTransformer(
clip=self.config["augmentor_clip_output"])
augmentor.to(self.device)
else:
raise ValueError("Unrecognized augmentor type: {}".format(
self.augmentor_type))
augmentor_optimizer = torch.optim.Adam(augmentor.parameters(), 3e-4)
augmentor_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
augmentor_optimizer,
T_max=len(train_loader),
eta_min=0,
last_epoch=-1)
optimizer = torch.optim.Adam(
list(model.parameters()),
3e-4,
weight_decay=eval(self.config["weight_decay"]),
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1)
if apex_support and self.config["fp16_precision"]:
model, optimizer = amp.initialize(model,
optimizer,
opt_level="O2",
keep_batchnorm_fp32=True)
model_checkpoints_folder = os.path.join(self.writer.log_dir,
"checkpoints")
# save config file
_save_config_file(model_checkpoints_folder)
n_iter = 0
valid_n_iter = 0
best_valid_loss = np.inf
for epoch_counter in range(self.config["epochs"]):
print("====== Epoch {} =======".format(epoch_counter))
for (xis, xjs), _ in train_loader:
optimizer.zero_grad()
xis = xis.to(self.device)
xjs = xjs.to(self.device)
loss = self._adv_step(model, augmentor, xis, xjs, n_iter)
if n_iter % self.config["log_every_n_steps"] == 0:
self.writer.add_scalar("train_loss",
loss,
global_step=n_iter)
if apex_support and self.config["fp16_precision"]:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
# for p in augmentor.parameters():
# # print(p.name)
# p.grad *= -1.0
optimizer.step()
# Update augmentor
augmentor_optimizer.zero_grad()
loss = self._adv_step(model, augmentor, xis, xjs, n_iter)
if self.augmentor_loss_type == "hinge":
loss = torch.clamp(loss, 0.0, 5.4)
loss *= -1.0
loss.backward()
augmentor_optimizer.step()
n_iter += 1
# validate the model if requested
if epoch_counter % self.config["eval_every_n_epochs"] == 0:
valid_loss = self._validate(model, augmentor, valid_loader)
if valid_loss < best_valid_loss:
# save the model weights
best_valid_loss = valid_loss
torch.save(
model.state_dict(),
os.path.join(model_checkpoints_folder, "model.pth"),
)
print("validation loss: ", valid_loss)
self.writer.add_scalar("validation_loss",
valid_loss,
global_step=valid_n_iter)
valid_n_iter += 1
# warmup for the first 10 epochs
if epoch_counter >= 10:
scheduler.step()
augmentor_scheduler.step()
self.writer.add_scalar("cosine_lr_decay",
scheduler.get_lr()[0],
global_step=n_iter)
def test(self):
train_loader, valid_loader, test_loader = self.dataset.get_data_loaders(
)
model = ResNetSimCLR(**self.config["model"]).to(self.device)
model = self._load_pre_trained_weights(model)
self.eval(test_loader, model)
test_loss = (test_loss * i + criterion(out, labels).item()) / (i + 1) res = torch.argmax(out, 1) correct += (res == labels).sum().item() tot += labels.size(0) return 100 - 100 * correct / tot, test_loss config_path = "./best_run/config.yaml" model_path = "./best_run/model.pth" data_path = "./data" config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) device = 'cuda:1' if torch.cuda.is_available() else 'cpu' model = ResNetSimCLR(**config["model"]).to(device) # x = torch.randn(2, 3, 32, 32).to(device) # print(model(x)) # exit() state_dict = torch.load(model_path) model.load_state_dict(state_dict) lr = 0.001 num_epoch = 200 batch_size = 512 num_classes = 10 transform = transforms.Compose( [transforms.ToTensor(),]) # may need to rewrite
