init
This commit is contained in:
@@ -0,0 +1,33 @@
|
||||
def setup(mode, P):
|
||||
fname = f'{P.dataset}_{P.model}_{mode}_{P.res}'
|
||||
|
||||
if mode == 'sup_linear':
|
||||
from .sup_linear import train
|
||||
elif mode == 'sup_CSI_linear':
|
||||
from .sup_CSI_linear import train
|
||||
elif mode == 'sup_simclr':
|
||||
from .sup_simclr import train
|
||||
elif mode == 'sup_simclr_CSI':
|
||||
assert P.batch_size == 32
|
||||
# currently only support rotation
|
||||
from .sup_simclr_CSI import train
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
if P.suffix is not None:
|
||||
fname += f'_{P.suffix}'
|
||||
|
||||
return train, fname
|
||||
|
||||
|
||||
def update_comp_loss(loss_dict, loss_in, loss_out, loss_diff, batch_size):
|
||||
loss_dict['pos'].update(loss_in, batch_size)
|
||||
loss_dict['neg'].update(loss_out, batch_size)
|
||||
loss_dict['diff'].update(loss_diff, batch_size)
|
||||
|
||||
|
||||
def summary_comp_loss(logger, tag, loss_dict, epoch):
|
||||
logger.scalar_summary(f'{tag}/pos', loss_dict['pos'].average, epoch)
|
||||
logger.scalar_summary(f'{tag}/neg', loss_dict['neg'].average, epoch)
|
||||
logger.scalar_summary(f'{tag}', loss_dict['diff'].average, epoch)
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,130 @@
|
||||
import time
|
||||
|
||||
import torch.optim
|
||||
import torch.optim.lr_scheduler as lr_scheduler
|
||||
|
||||
import models.transform_layers as TL
|
||||
from utils.utils import AverageMeter, normalize
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
hflip = TL.HorizontalFlipLayer().to(device)
|
||||
|
||||
|
||||
def train(P, epoch, model, criterion, optimizer, scheduler, loader, logger=None,
|
||||
simclr_aug=None, linear=None, linear_optim=None):
|
||||
|
||||
if P.multi_gpu:
|
||||
rotation_linear = model.module.shift_cls_layer
|
||||
joint_linear = model.module.joint_distribution_layer
|
||||
else:
|
||||
rotation_linear = model.shift_cls_layer
|
||||
joint_linear = model.joint_distribution_layer
|
||||
|
||||
if epoch == 1:
|
||||
# define optimizer and save in P (argument)
|
||||
milestones = [int(0.6 * P.epochs), int(0.75 * P.epochs), int(0.9 * P.epochs)]
|
||||
|
||||
linear_optim = torch.optim.SGD(linear.parameters(),
|
||||
lr=1e-1, weight_decay=P.weight_decay)
|
||||
P.linear_optim = linear_optim
|
||||
P.linear_scheduler = lr_scheduler.MultiStepLR(P.linear_optim, gamma=0.1, milestones=milestones)
|
||||
|
||||
rotation_linear_optim = torch.optim.SGD(rotation_linear.parameters(),
|
||||
lr=1e-1, weight_decay=P.weight_decay)
|
||||
P.rotation_linear_optim = rotation_linear_optim
|
||||
P.rot_scheduler = lr_scheduler.MultiStepLR(P.rotation_linear_optim, gamma=0.1, milestones=milestones)
|
||||
|
||||
joint_linear_optim = torch.optim.SGD(joint_linear.parameters(),
|
||||
lr=1e-1, weight_decay=P.weight_decay)
|
||||
P.joint_linear_optim = joint_linear_optim
|
||||
P.joint_scheduler = lr_scheduler.MultiStepLR(P.joint_linear_optim, gamma=0.1, milestones=milestones)
|
||||
|
||||
if logger is None:
|
||||
log_ = print
|
||||
else:
|
||||
log_ = logger.log
|
||||
|
||||
batch_time = AverageMeter()
|
||||
data_time = AverageMeter()
|
||||
|
||||
losses = dict()
|
||||
losses['cls'] = AverageMeter()
|
||||
losses['rot'] = AverageMeter()
|
||||
|
||||
check = time.time()
|
||||
for n, (images, labels) in enumerate(loader):
|
||||
model.eval()
|
||||
count = n * P.n_gpus # number of trained samples
|
||||
|
||||
data_time.update(time.time() - check)
|
||||
check = time.time()
|
||||
|
||||
### SimCLR loss ###
|
||||
if P.dataset != 'imagenet':
|
||||
batch_size = images.size(0)
|
||||
images = images.to(device)
|
||||
images = hflip(images) # 2B with hflip
|
||||
else:
|
||||
batch_size = images[0].size(0)
|
||||
images = images[0].to(device)
|
||||
|
||||
labels = labels.to(device)
|
||||
images = torch.cat([torch.rot90(images, rot, (2, 3)) for rot in range(4)]) # 4B
|
||||
rot_labels = torch.cat([torch.ones_like(labels) * k for k in range(4)], 0) # B -> 4B
|
||||
joint_labels = torch.cat([labels + P.n_classes * i for i in range(4)], dim=0)
|
||||
|
||||
images = simclr_aug(images) # simclr augmentation
|
||||
_, outputs_aux = model(images, penultimate=True)
|
||||
penultimate = outputs_aux['penultimate'].detach()
|
||||
|
||||
outputs = linear(penultimate[0:batch_size]) # only use 0 degree samples for linear eval
|
||||
outputs_rot = rotation_linear(penultimate)
|
||||
outputs_joint = joint_linear(penultimate)
|
||||
|
||||
loss_ce = criterion(outputs, labels)
|
||||
loss_rot = criterion(outputs_rot, rot_labels)
|
||||
loss_joint = criterion(outputs_joint, joint_labels)
|
||||
|
||||
### CE loss ###
|
||||
P.linear_optim.zero_grad()
|
||||
loss_ce.backward()
|
||||
P.linear_optim.step()
|
||||
|
||||
### Rot loss ###
|
||||
P.rotation_linear_optim.zero_grad()
|
||||
loss_rot.backward()
|
||||
P.rotation_linear_optim.step()
|
||||
|
||||
### Joint loss ###
|
||||
P.joint_linear_optim.zero_grad()
|
||||
loss_joint.backward()
|
||||
P.joint_linear_optim.step()
|
||||
|
||||
### optimizer learning rate ###
|
||||
lr = P.linear_optim.param_groups[0]['lr']
|
||||
|
||||
batch_time.update(time.time() - check)
|
||||
|
||||
### Log losses ###
|
||||
losses['cls'].update(loss_ce.item(), batch_size)
|
||||
losses['rot'].update(loss_rot.item(), batch_size)
|
||||
|
||||
if count % 50 == 0:
|
||||
log_('[Epoch %3d; %3d] [Time %.3f] [Data %.3f] [LR %.5f]\n'
|
||||
'[LossC %f] [LossR %f]' %
|
||||
(epoch, count, batch_time.value, data_time.value, lr,
|
||||
losses['cls'].value, losses['rot'].value))
|
||||
check = time.time()
|
||||
|
||||
P.linear_scheduler.step()
|
||||
P.rot_scheduler.step()
|
||||
P.joint_scheduler.step()
|
||||
|
||||
log_('[DONE] [Time %.3f] [Data %.3f] [LossC %f] [LossR %f]' %
|
||||
(batch_time.average, data_time.average,
|
||||
losses['cls'].average, losses['rot'].average))
|
||||
|
||||
if logger is not None:
|
||||
logger.scalar_summary('train/loss_cls', losses['cls'].average, epoch)
|
||||
logger.scalar_summary('train/loss_rot', losses['rot'].average, epoch)
|
||||
logger.scalar_summary('train/batch_time', batch_time.average, epoch)
|
||||
@@ -0,0 +1,91 @@
|
||||
import time
|
||||
|
||||
import torch.optim
|
||||
import torch.optim.lr_scheduler as lr_scheduler
|
||||
|
||||
import models.transform_layers as TL
|
||||
from utils.utils import AverageMeter, normalize
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
hflip = TL.HorizontalFlipLayer().to(device)
|
||||
|
||||
|
||||
def train(P, epoch, model, criterion, optimizer, scheduler, loader, logger=None,
|
||||
simclr_aug=None, linear=None, linear_optim=None):
|
||||
|
||||
if epoch == 1:
|
||||
# define optimizer and save in P (argument)
|
||||
milestones = [int(0.6 * P.epochs), int(0.75 * P.epochs), int(0.9 * P.epochs)]
|
||||
|
||||
linear_optim = torch.optim.SGD(linear.parameters(),
|
||||
lr=1e-1, weight_decay=P.weight_decay)
|
||||
P.linear_optim = linear_optim
|
||||
P.linear_scheduler = lr_scheduler.MultiStepLR(P.linear_optim, gamma=0.1, milestones=milestones)
|
||||
|
||||
if logger is None:
|
||||
log_ = print
|
||||
else:
|
||||
log_ = logger.log
|
||||
|
||||
batch_time = AverageMeter()
|
||||
data_time = AverageMeter()
|
||||
|
||||
losses = dict()
|
||||
losses['cls'] = AverageMeter()
|
||||
|
||||
check = time.time()
|
||||
for n, (images, labels) in enumerate(loader):
|
||||
model.eval()
|
||||
count = n * P.n_gpus # number of trained samples
|
||||
|
||||
data_time.update(time.time() - check)
|
||||
check = time.time()
|
||||
|
||||
### SimCLR loss ###
|
||||
if P.dataset != 'imagenet':
|
||||
batch_size = images.size(0)
|
||||
images = images.to(device)
|
||||
images = hflip(images) # 2B with hflip
|
||||
else:
|
||||
batch_size = images[0].size(0)
|
||||
images = images[0].to(device)
|
||||
|
||||
labels = labels.to(device)
|
||||
|
||||
images = simclr_aug(images) # simclr augmentation
|
||||
_, outputs_aux = model(images, penultimate=True)
|
||||
penultimate = outputs_aux['penultimate'].detach()
|
||||
|
||||
outputs = linear(penultimate[0:batch_size]) # only use 0 degree samples for linear eval
|
||||
|
||||
loss_ce = criterion(outputs, labels)
|
||||
|
||||
### CE loss ###
|
||||
P.linear_optim.zero_grad()
|
||||
loss_ce.backward()
|
||||
P.linear_optim.step()
|
||||
|
||||
### optimizer learning rate ###
|
||||
lr = P.linear_optim.param_groups[0]['lr']
|
||||
|
||||
batch_time.update(time.time() - check)
|
||||
|
||||
### Log losses ###
|
||||
losses['cls'].update(loss_ce.item(), batch_size)
|
||||
|
||||
if count % 50 == 0:
|
||||
log_('[Epoch %3d; %3d] [Time %.3f] [Data %.3f] [LR %.5f]\n'
|
||||
'[LossC %f]' %
|
||||
(epoch, count, batch_time.value, data_time.value, lr,
|
||||
losses['cls'].value, ))
|
||||
check = time.time()
|
||||
|
||||
P.linear_scheduler.step()
|
||||
|
||||
log_('[DONE] [Time %.3f] [Data %.3f] [LossC %f]' %
|
||||
(batch_time.average, data_time.average,
|
||||
losses['cls'].average))
|
||||
|
||||
if logger is not None:
|
||||
logger.scalar_summary('train/loss_cls', losses['cls'].average, epoch)
|
||||
logger.scalar_summary('train/batch_time', batch_time.average, epoch)
|
||||
@@ -0,0 +1,104 @@
|
||||
import time
|
||||
|
||||
import torch.optim
|
||||
|
||||
import models.transform_layers as TL
|
||||
from training.contrastive_loss import get_similarity_matrix, Supervised_NT_xent
|
||||
from utils.utils import AverageMeter, normalize
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
hflip = TL.HorizontalFlipLayer().to(device)
|
||||
|
||||
|
||||
def train(P, epoch, model, criterion, optimizer, scheduler, loader, logger=None,
|
||||
simclr_aug=None, linear=None, linear_optim=None):
|
||||
|
||||
assert simclr_aug is not None
|
||||
assert P.sim_lambda == 1.0
|
||||
|
||||
if logger is None:
|
||||
log_ = print
|
||||
else:
|
||||
log_ = logger.log
|
||||
|
||||
batch_time = AverageMeter()
|
||||
data_time = AverageMeter()
|
||||
|
||||
losses = dict()
|
||||
losses['cls'] = AverageMeter()
|
||||
losses['sim'] = AverageMeter()
|
||||
losses['simnorm'] = AverageMeter()
|
||||
|
||||
check = time.time()
|
||||
for n, (images, labels) in enumerate(loader):
|
||||
model.train()
|
||||
count = n * P.n_gpus # number of trained samples
|
||||
|
||||
data_time.update(time.time() - check)
|
||||
check = time.time()
|
||||
|
||||
### SimCLR loss ###
|
||||
if P.dataset != 'imagenet' and P.dataset != 'CNMC' and P.dataset != 'CNMC_grayscale':
|
||||
batch_size = images.size(0)
|
||||
images = images.to(device)
|
||||
images_pair = hflip(images.repeat(2, 1, 1, 1)) # 2B with hflip
|
||||
else:
|
||||
batch_size = images[0].size(0)
|
||||
images1, images2 = images[0].to(device), images[1].to(device)
|
||||
images_pair = torch.cat([images1, images2], dim=0) # 2B
|
||||
|
||||
labels = labels.to(device)
|
||||
|
||||
images_pair = simclr_aug(images_pair) # simclr augmentation
|
||||
|
||||
_, outputs_aux = model(images_pair, simclr=True, penultimate=True)
|
||||
|
||||
simclr = normalize(outputs_aux['simclr']) # normalize
|
||||
sim_matrix = get_similarity_matrix(simclr, multi_gpu=P.multi_gpu)
|
||||
loss_sim = Supervised_NT_xent(sim_matrix, labels=labels, temperature=0.07, multi_gpu=P.multi_gpu) * P.sim_lambda
|
||||
|
||||
### total loss ###
|
||||
loss = loss_sim
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
scheduler.step(epoch - 1 + n / len(loader))
|
||||
lr = optimizer.param_groups[0]['lr']
|
||||
|
||||
batch_time.update(time.time() - check)
|
||||
|
||||
### Post-processing stuffs ###
|
||||
simclr_norm = outputs_aux['simclr'].norm(dim=1).mean()
|
||||
|
||||
### Linear evaluation ###
|
||||
outputs_linear_eval = linear(outputs_aux['penultimate'].detach())
|
||||
loss_linear = criterion(outputs_linear_eval, labels.repeat(2))
|
||||
|
||||
linear_optim.zero_grad()
|
||||
loss_linear.backward()
|
||||
linear_optim.step()
|
||||
|
||||
### Log losses ###
|
||||
losses['cls'].update(0, batch_size)
|
||||
losses['sim'].update(loss_sim.item(), batch_size)
|
||||
losses['simnorm'].update(simclr_norm.item(), batch_size)
|
||||
|
||||
if count % 50 == 0:
|
||||
log_('[Epoch %3d; %3d] [Time %.3f] [Data %.3f] [LR %.5f]\n'
|
||||
'[LossC %f] [LossSim %f] [SimNorm %f]' %
|
||||
(epoch, count, batch_time.value, data_time.value, lr,
|
||||
losses['cls'].value, losses['sim'].value, losses['simnorm'].value))
|
||||
|
||||
check = time.time()
|
||||
|
||||
log_('[DONE] [Time %.3f] [Data %.3f] [LossC %f] [LossSim %f] [SimNorm %f]' %
|
||||
(batch_time.average, data_time.average,
|
||||
losses['cls'].average, losses['sim'].average, losses['simnorm'].average))
|
||||
|
||||
if logger is not None:
|
||||
logger.scalar_summary('train/loss_cls', losses['cls'].average, epoch)
|
||||
logger.scalar_summary('train/loss_sim', losses['sim'].average, epoch)
|
||||
logger.scalar_summary('train/batch_time', batch_time.average, epoch)
|
||||
logger.scalar_summary('train/simclr_norm', losses['simnorm'].average, epoch)
|
||||
@@ -0,0 +1,111 @@
|
||||
import time
|
||||
|
||||
import torch.optim
|
||||
|
||||
import models.transform_layers as TL
|
||||
from training.contrastive_loss import get_similarity_matrix, Supervised_NT_xent
|
||||
from utils.utils import AverageMeter, normalize
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
hflip = TL.HorizontalFlipLayer().to(device)
|
||||
|
||||
|
||||
def train(P, epoch, model, criterion, optimizer, scheduler, loader, logger=None,
|
||||
simclr_aug=None, linear=None, linear_optim=None):
|
||||
|
||||
# currently only support rotation shifting augmentation
|
||||
assert simclr_aug is not None
|
||||
assert P.sim_lambda == 1.0
|
||||
|
||||
if logger is None:
|
||||
log_ = print
|
||||
else:
|
||||
log_ = logger.log
|
||||
|
||||
batch_time = AverageMeter()
|
||||
data_time = AverageMeter()
|
||||
|
||||
losses = dict()
|
||||
losses['cls'] = AverageMeter()
|
||||
losses['sim'] = AverageMeter()
|
||||
|
||||
check = time.time()
|
||||
for n, (images, labels) in enumerate(loader):
|
||||
model.train()
|
||||
count = n * P.n_gpus # number of trained samples
|
||||
|
||||
data_time.update(time.time() - check)
|
||||
check = time.time()
|
||||
|
||||
### SimCLR loss ###
|
||||
if P.dataset != 'imagenet' and P.dataset != 'CNMC' and P.dataset != 'CNMC_grayscale':
|
||||
batch_size = images.size(0)
|
||||
images = images.to(device)
|
||||
images1, images2 = hflip(images.repeat(2, 1, 1, 1)).chunk(2) # hflip
|
||||
else:
|
||||
batch_size = images[0].size(0)
|
||||
images1, images2 = images[0].to(device), images[1].to(device)
|
||||
#print("\nImages" + str(images.shape) + "\n")
|
||||
|
||||
images1 = torch.cat([torch.rot90(images1, rot, (2, 3)) for rot in range(4)]) # 4B
|
||||
images2 = torch.cat([torch.rot90(images2, rot, (2, 3)) for rot in range(4)]) # 4B
|
||||
images_pair = torch.cat([images1, images2], dim=0) # 8B
|
||||
|
||||
labels = labels.to(device)
|
||||
rot_sim_labels = torch.cat([labels + P.n_classes * i for i in range(4)], dim=0)
|
||||
rot_sim_labels = rot_sim_labels.to(device)
|
||||
|
||||
images_pair = simclr_aug(images_pair) # simclr augment
|
||||
_, outputs_aux = model(images_pair, simclr=True, penultimate=True)
|
||||
|
||||
simclr = normalize(outputs_aux['simclr']) # normalize
|
||||
sim_matrix = get_similarity_matrix(simclr, multi_gpu=P.multi_gpu)
|
||||
loss_sim = Supervised_NT_xent(sim_matrix, labels=rot_sim_labels,
|
||||
temperature=0.07, multi_gpu=P.multi_gpu) * P.sim_lambda
|
||||
|
||||
### total loss ###
|
||||
loss = loss_sim
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
scheduler.step(epoch - 1 + n / len(loader))
|
||||
lr = optimizer.param_groups[0]['lr']
|
||||
|
||||
batch_time.update(time.time() - check)
|
||||
|
||||
### Post-processing stuffs ###
|
||||
penul_1 = outputs_aux['penultimate'][:batch_size]
|
||||
penul_2 = outputs_aux['penultimate'][4 * batch_size: 5 * batch_size]
|
||||
outputs_aux['penultimate'] = torch.cat([penul_1, penul_2]) # only use original rotation
|
||||
|
||||
### Linear evaluation ###
|
||||
outputs_linear_eval = linear(outputs_aux['penultimate'].detach())
|
||||
loss_linear = criterion(outputs_linear_eval, labels.repeat(2))
|
||||
|
||||
linear_optim.zero_grad()
|
||||
loss_linear.backward()
|
||||
linear_optim.step()
|
||||
|
||||
### Log losses ###
|
||||
losses['cls'].update(0, batch_size)
|
||||
losses['sim'].update(loss_sim.item(), batch_size)
|
||||
|
||||
if count % 50 == 0:
|
||||
log_('[Epoch %3d; %3d] [Time %.3f] [Data %.3f] [LR %.5f]\n'
|
||||
'[LossC %f] [LossSim %f]' %
|
||||
(epoch, count, batch_time.value, data_time.value, lr,
|
||||
losses['cls'].value, losses['sim'].value))
|
||||
|
||||
check = time.time()
|
||||
|
||||
log_('[DONE] [Time %.3f] [Data %.3f] [LossC %f] [LossSim %f]' %
|
||||
(batch_time.average, data_time.average,
|
||||
losses['cls'].average, losses['sim'].average))
|
||||
|
||||
if logger is not None:
|
||||
logger.scalar_summary('train/loss_cls', losses['cls'].average, epoch)
|
||||
logger.scalar_summary('train/loss_sim', losses['sim'].average, epoch)
|
||||
logger.scalar_summary('train/batch_time', batch_time.average, epoch)
|
||||
|
||||
Reference in New Issue
Block a user