def validate(val_loader, model, criterion, log=None, use_cuda=True):
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
for i, (input, target) in enumerate(val_loader):
if use_cuda:
target = target.cuda()
input = input.cuda()
with torch.no_grad():
# compute output
output = model(input)
# measure accuracy and record loss
prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
top1.update(prec1.item(), input.size(0))
top5.update(prec5.item(), input.size(0))
if log:
print_log(' **Test** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1,
top5=top5,
error1=100-top1.avg), log)
else:
print(' **Test** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, top5=top5, error1=100-top1.avg))
return top1.avg
def save_checkpoint(self, val_loss, model, optimizer, log):
'''Saves model when validation loss decrease.'''
if self.verbose:
print_log((
f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...'
), log)
state = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(state, self.save_name)
self.val_loss_min = val_loss
def train_model(data_loader, model, criterion, optimizer, epoch, log,
print_freq=200, use_cuda=True):
# train function (forward, backward, update)
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
# switch to train mode
model.train()
end = time.time()
for iteration, (input, target) in enumerate(data_loader):
# measure data loading time
data_time.update(time.time() - end)
if use_cuda:
target = target.cuda()
input = input.cuda()
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
if len(target.shape) > 1:
target = torch.argmax(target, dim=-1)
prec1, = accuracy(output.data, target, topk=(1,))
losses.update(loss.item(), input.size(0))
top1.update(prec1.item(), input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if iteration % print_freq == 0:
print_log(' Epoch: [{:03d}][{:03d}/{:03d}] '
'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) '
'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
'Loss {loss.val:.4f} ({loss.avg:.4f}) '
'Prec@1 {top1.val:.3f} ({top1.avg:.3f}) '.format(
epoch, iteration, len(data_loader), batch_time=batch_time,
data_time=data_time, loss=losses, top1=top1) + time_string(), log)
print_log(' **Train** Prec@1 {top1.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1, error1=100-top1.avg), log)
return top1.avg, losses.avg
def test(args, model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
utils.print_log('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)), args)
def train(args, model, device, train_loader, optimizer, epoch):
"""
Training the classifier.
"""
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
utils.print_log('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()), args)
if args.dry_run:
break
def __call__(self, val_loss, model, optimizer, log):
score = -val_loss
if self.best_score is None:
self.best_score = score
self.save_checkpoint(val_loss, model, optimizer, log)
elif score < self.best_score - self.delta:
self.counter += 1
print_log((
f'EarlyStopping counter: {self.counter} out of {self.patience}'
), log)
if self.counter >= self.patience:
self.early_stop = True
else:
self.best_score = score
self.save_checkpoint(val_loss, model, optimizer, log)
self.counter = 0
return self.early_stop
# utils.load_state_dict(model, weight_file)
# model.fc.reset_parameters()
criterion = nn.CrossEntropyLoss()
if cuda:
model = model.cuda()
criterion = criterion.cuda()
optim = torch.optim.SGD(model.parameters(),
lr=train_cfg['lr'],
momentum=train_cfg['momentum'],
weight_decay=train_cfg['weight_decay'])
# lr_scheduler = torch.optim.lr_scheduler.StepLR(optim, train_cfg['step_size'], gamma=train_cfg['gamma'], last_epoch=-1)
lr_scheduler = None
log = utils.print_log(configure['log_dir'], [net_cfg['type'], timestamp])
log.write(str(net_cfg))
log.write(str(train_cfg))
epoch_time = utils.AverageMeter()
batch_time = utils.AverageMeter()
data_time = utils.AverageMeter()
losses = utils.AverageMeter()
top1 = utils.AverageMeter()
# --------------------train & validation & save checkpoint---------------- #
epoch = 0
last_iteration = 0
print_freq = 1
best_top1 = 0
max_epoch = train_cfg['max_epoch']
print("train max epoch {0}".format(max_epoch))
def train():
train_data = ACNet_data.SUNRGBD(transform=transforms.Compose([
ACNet_data.scaleNorm(),
ACNet_data.RandomScale((1.0, 1.4)),
ACNet_data.RandomHSV((0.9, 1.1), (0.9, 1.1), (25, 25)),
ACNet_data.RandomCrop(image_h, image_w),
ACNet_data.RandomFlip(),
ACNet_data.ToTensor(),
ACNet_data.Normalize()
]),
phase_train=True,
data_dir=args.data_dir)
train_loader = DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=False)
num_train = len(train_data)
if args.last_ckpt:
model = ACNet_models_V1.ACNet(num_class=40, pretrained=False)
else:
model = ACNet_models_V1.ACNet(num_class=40, pretrained=True)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
# CEL_weighted = utils.CrossEntropyLoss2d()
CEL_weighted = utils.FocalLoss2d(weight=nyuv2_frq, gamma=2)
model.train()
model.to(device)
CEL_weighted.to(device)
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
global_step = 0
if args.last_ckpt:
global_step, args.start_epoch = load_ckpt(model, optimizer,
args.last_ckpt, device)
#hxx for finetuing
# lr_decay_lambda = lambda epoch: 0.2 * args.lr_decay_rate ** ((epoch - args.start_epoch) // args.lr_epoch_per_decay)
lr_decay_lambda = lambda epoch: (1 - (epoch - args.start_epoch) /
(args.epochs - args.start_epoch))**0.9
scheduler = LambdaLR(optimizer, lr_lambda=lr_decay_lambda)
writer = SummaryWriter(args.summary_dir)
for epoch in range(int(args.start_epoch), args.epochs):
# if (epoch - args.start_epoch) % args.lr_epoch_per_decay == 0:
scheduler.step(epoch)
local_count = 0
last_count = 0
end_time = time.time()
if epoch % args.save_epoch_freq == 0 and epoch != args.start_epoch:
save_ckpt(args.ckpt_dir, model, optimizer, global_step, epoch,
local_count, num_train)
for batch_idx, sample in enumerate(train_loader):
image = sample['image'].to(device)
depth = sample['depth'].to(device)
target_scales = [
sample[s].to(device)
for s in ['label', 'label2', 'label3', 'label4', 'label5']
]
optimizer.zero_grad()
pred_scales = model(image, depth, args.checkpoint)
loss = CEL_weighted(pred_scales, target_scales)
loss.backward()
optimizer.step()
local_count += image.data.shape[0]
global_step += 1
if global_step % args.print_freq == 0 or global_step == 1:
time_inter = time.time() - end_time
count_inter = local_count - last_count
print_log(global_step, epoch, local_count, count_inter,
num_train, loss, time_inter)
end_time = time.time()
last_count = local_count
save_ckpt(args.ckpt_dir, model, optimizer, global_step, args.epochs, 0,
num_train)
print("Training completed ")
def main():
args = parse_arguments()
random.seed(args.pretrained_seed)
torch.manual_seed(args.pretrained_seed)
if args.use_cuda:
torch.cuda.manual_seed_all(args.pretrained_seed)
cudnn.benchmark = True
# get the result path to store the results
result_path = get_result_path(dataset_name=args.dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed,
result_subfolder=args.result_subfolder,
postfix=args.postfix)
# Init logger
log_file_name = os.path.join(result_path, 'log.txt')
print("Log file: {}".format(log_file_name))
log = open(log_file_name, 'w')
print_log('save path : {}'.format(result_path), log)
state = {k: v for k, v in args._get_kwargs()}
for key, value in state.items():
print_log("{} : {}".format(key, value), log)
print_log("Random Seed: {}".format(args.pretrained_seed), log)
print_log("Python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("Torch version : {}".format(torch.__version__), log)
print_log("Cudnn version : {}".format(torch.backends.cudnn.version()), log)
_, pretrained_data_test = get_data(args.pretrained_dataset, args.pretrained_dataset)
pretrained_data_test_loader = torch.utils.data.DataLoader(pretrained_data_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
##### Dataloader for training ####
num_classes, (mean, std), input_size, num_channels = get_data_specs(args.pretrained_dataset)
data_train, _ = get_data(args.dataset, args.pretrained_dataset)
data_train_loader = torch.utils.data.DataLoader(data_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
####################################
# Init model, criterion, and optimizer
print_log("=> Creating model '{}'".format(args.pretrained_arch), log)
# get a path for loading the model to be attacked
model_path = get_model_path(dataset_name=args.pretrained_dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed)
model_weights_path = os.path.join(model_path, "checkpoint.pth.tar")
target_network = get_network(args.pretrained_arch,
input_size=input_size,
num_classes=num_classes,
finetune=False)
print_log("=> Network :\n {}".format(target_network), log)
target_network = torch.nn.DataParallel(target_network, device_ids=list(range(args.ngpu)))
# Set the target model into evaluation mode
target_network.eval()
# Imagenet models use the pretrained pytorch weights
if args.pretrained_dataset != "imagenet":
network_data = torch.load(model_weights_path)
target_network.load_state_dict(network_data['state_dict'])
# Set all weights to not trainable
set_parameter_requires_grad(target_network, requires_grad=False)
non_trainale_params = get_num_non_trainable_parameters(target_network)
trainale_params = get_num_trainable_parameters(target_network)
total_params = get_num_parameters(target_network)
print_log("Target Network Trainable parameters: {}".format(trainale_params), log)
print_log("Target Network Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Target Network Total # parameters: {}".format(total_params), log)
print_log("=> Inserting Generator", log)
generator = UAP(shape=(input_size, input_size),
num_channels=num_channels,
mean=mean,
std=std,
use_cuda=args.use_cuda)
print_log("=> Generator :\n {}".format(generator), log)
non_trainale_params = get_num_non_trainable_parameters(generator)
trainale_params = get_num_trainable_parameters(generator)
total_params = get_num_parameters(generator)
print_log("Generator Trainable parameters: {}".format(trainale_params), log)
print_log("Generator Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Generator Total # parameters: {}".format(total_params), log)
perturbed_net = nn.Sequential(OrderedDict([('generator', generator), ('target_model', target_network)]))
perturbed_net = torch.nn.DataParallel(perturbed_net, device_ids=list(range(args.ngpu)))
non_trainale_params = get_num_non_trainable_parameters(perturbed_net)
trainale_params = get_num_trainable_parameters(perturbed_net)
total_params = get_num_parameters(perturbed_net)
print_log("Perturbed Net Trainable parameters: {}".format(trainale_params), log)
print_log("Perturbed Net Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Perturbed Net Total # parameters: {}".format(total_params), log)
# Set the target model into evaluation mode
perturbed_net.module.target_model.eval()
perturbed_net.module.generator.train()
if args.loss_function == "ce":
criterion = torch.nn.CrossEntropyLoss()
elif args.loss_function == "neg_ce":
criterion = NegativeCrossEntropy()
elif args.loss_function == "logit":
criterion = LogitLoss(num_classes=num_classes, use_cuda=args.use_cuda)
elif args.loss_function == "bounded_logit":
criterion = BoundedLogitLoss(num_classes=num_classes, confidence=args.confidence, use_cuda=args.use_cuda)
elif args.loss_function == "bounded_logit_fixed_ref":
criterion = BoundedLogitLossFixedRef(num_classes=num_classes, confidence=args.confidence, use_cuda=args.use_cuda)
elif args.loss_function == "bounded_logit_neg":
criterion = BoundedLogitLoss_neg(num_classes=num_classes, confidence=args.confidence, use_cuda=args.use_cuda)
else:
raise ValueError
if args.use_cuda:
target_network.cuda()
generator.cuda()
perturbed_net.cuda()
criterion.cuda()
optimizer = torch.optim.Adam(perturbed_net.parameters(), lr=state['learning_rate'])
# Measure the time needed for the UAP generation
start = time.time()
train(data_loader=data_train_loader,
model=perturbed_net,
criterion=criterion,
optimizer=optimizer,
epsilon=args.epsilon,
num_iterations=args.num_iterations,
targeted=args.targeted,
target_class=args.target_class,
log=log,
print_freq=args.print_freq,
use_cuda=args.use_cuda)
end = time.time()
print_log("Time needed for UAP generation: {}".format(end - start), log)
# evaluate
print_log("Final evaluation:", log)
metrics_evaluate(data_loader=pretrained_data_test_loader,
target_model=target_network,
perturbed_model=perturbed_net,
targeted=args.targeted,
target_class=args.target_class,
log=log,
use_cuda=args.use_cuda)
save_checkpoint({
'arch' : args.pretrained_arch,
# 'state_dict' : perturbed_net.state_dict(),
'state_dict' : perturbed_net.module.generator.state_dict(),
'optimizer' : optimizer.state_dict(),
'args' : copy.deepcopy(args),
}, result_path, 'checkpoint.pth.tar')
log.close()
def main():
args = parse_arguments()
random.seed(args.pretrained_seed)
torch.manual_seed(args.pretrained_seed)
if args.use_cuda:
torch.cuda.manual_seed_all(args.pretrained_seed)
cudnn.benchmark = True
# Get data specs
num_classes, (mean, std), input_size, num_channels = get_data_specs(args.pretrained_dataset, args.pretrained_arch)
# Construct the array other classes:
if args.pretrained_dataset in ["imagenet", "ycb"]:
other_classes = args.source_classes
else:
all_classes = np.arange(num_classes)
other_classes = [int(cl) for cl in all_classes if cl not in args.source_classes]
half_batch_size = args.batch_size//2
# get the result path to store the results
result_path = get_result_path(dataset_name=args.pretrained_dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed,
result_subfolder=args.result_subfolder,
source_class=args.source_classes,
sink_class=args.sink_classes,
postfix=args.postfix)
# Init logger
log_file_name = os.path.join(result_path, 'log.txt')
print("Log file: {}".format(log_file_name))
log = open(log_file_name, 'w')
print_log('save path : {}'.format(result_path), log)
state = {k: v for k, v in args._get_kwargs()}
for key, value in state.items():
print_log("{} : {}".format(key, value), log)
print_log("Random Seed: {}".format(args.pretrained_seed), log)
print_log("Python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("Torch version : {}".format(torch.__version__), log)
print_log("Cudnn version : {}".format(torch.backends.cudnn.version()), log)
data_train_sources, data_test_sources = get_data(args.pretrained_dataset,
mean=mean,
std=std,
input_size=input_size,
classes=args.source_classes,
train_samples_per_class=args.num_train_samples_per_class)
data_train_sources_loader = torch.utils.data.DataLoader(data_train_sources,
batch_size=half_batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
data_test_sources_loader = torch.utils.data.DataLoader(data_test_sources,
batch_size=half_batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
data_train_others, data_test_others = get_data(args.pretrained_dataset,
mean=mean,
std=std,
input_size=input_size,
classes=other_classes,
others=True,
train_samples_per_class=args.num_train_samples_per_class)
data_train_others_loader = torch.utils.data.DataLoader(data_train_others,
batch_size=half_batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
data_test_others_loader = torch.utils.data.DataLoader(data_test_others,
batch_size=half_batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# Init model, criterion, and optimizer
print_log("=> Creating model '{}'".format(args.pretrained_arch), log)
# get a path for loading the model to be attacked
model_path = get_model_path(dataset_name=args.pretrained_dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed)
model_weights_path = os.path.join(model_path, "checkpoint.pth.tar")
target_network = get_network(args.pretrained_arch, input_size=input_size, num_classes=num_classes, finetune=args.finetune)
# print_log("=> Network :\n {}".format(target_network), log)
target_network = torch.nn.DataParallel(target_network, device_ids=list(range(args.ngpu)))
# Set the target model into evaluation mode
target_network.eval()
# Imagenet models use the pretrained pytorch weights
if args.pretrained_dataset != "imagenet":
network_data = torch.load(model_weights_path)
target_network.load_state_dict(network_data['state_dict'])
# Set all weights to not trainable
set_parameter_requires_grad(target_network, requires_grad=False)
non_trainale_params = get_num_non_trainable_parameters(target_network)
trainale_params = get_num_trainable_parameters(target_network)
total_params = get_num_parameters(target_network)
print_log("Target Network Trainable parameters: {}".format(trainale_params), log)
print_log("Target Network Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Target Network Total # parameters: {}".format(total_params), log)
print_log("=> Inserting Generator", log)
generator = UAP(shape=(input_size, input_size),
num_channels=num_channels,
mean=mean,
std=std,
use_cuda=args.use_cuda)
print_log("=> Generator :\n {}".format(generator), log)
non_trainale_params = get_num_non_trainable_parameters(generator)
trainale_params = get_num_trainable_parameters(generator)
total_params = get_num_parameters(generator)
print_log("Generator Trainable parameters: {}".format(trainale_params), log)
print_log("Generator Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Generator Total # parameters: {}".format(total_params), log)
perturbed_net = nn.Sequential(OrderedDict([('generator', generator), ('target_model', target_network)]))
perturbed_net = torch.nn.DataParallel(perturbed_net, device_ids=list(range(args.ngpu)))
non_trainale_params = get_num_non_trainable_parameters(perturbed_net)
trainale_params = get_num_trainable_parameters(perturbed_net)
total_params = get_num_parameters(perturbed_net)
print_log("Perturbed Net Trainable parameters: {}".format(trainale_params), log)
print_log("Perturbed Net Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Perturbed Net Total # parameters: {}".format(total_params), log)
# Set the target model into evaluation mode
perturbed_net.module.target_model.eval()
perturbed_net.module.generator.train()
criterion = LossConstructor(source_classes=args.source_classes,
sink_classes=args.sink_classes,
num_classes=num_classes,
source_loss=args.source_loss,
others_loss=args.others_loss,
confidence=args.confidence,
alpha=args.alpha,
use_cuda=args.use_cuda)
if args.use_cuda:
target_network.cuda()
generator.cuda()
perturbed_net.cuda()
criterion.cuda()
optimizer = torch.optim.Adam(perturbed_net.parameters(),
lr=state['learning_rate']) # betas=(0.5, 0.999)
if args.pretrained_dataset not in ["imagenet", "ycb"]:
metrics_evaluate(source_loader=data_test_sources_loader,
others_loader=data_test_others_loader,
target_model=target_network,
perturbed_model=perturbed_net,
source_classes=args.source_classes,
sink_classes=args.sink_classes,
log=log,
use_cuda=args.use_cuda)
start_time = time.time()
train_half_half(sources_data_loader=data_train_sources_loader,
others_data_loader=data_train_others_loader,
model=perturbed_net,
target_model=target_network,
criterion=criterion,
optimizer=optimizer,
epsilon=args.epsilon,
num_iterations=args.num_iterations,
log=log,
print_freq=args.print_freq,
use_cuda=args.use_cuda)
end_time = time.time()
print_log("Elapsed generation time: {}".format(end_time-start_time), log)
# evaluate
print_log("Final evaluation:", log)
metrics_evaluate(source_loader=data_test_sources_loader,
others_loader=data_test_others_loader,
target_model=target_network,
perturbed_model=perturbed_net,
source_classes=args.source_classes,
sink_classes=args.sink_classes,
log=log,
use_cuda=args.use_cuda)
save_checkpoint({
'arch' : args.pretrained_arch,
'state_dict' : perturbed_net.state_dict(),
'optimizer' : optimizer.state_dict(),
'args' : copy.deepcopy(args),
}, result_path, 'checkpoint.pth.tar')
# Plot the adversarial perturbation
uap_numpy = perturbed_net.module.generator.uap.detach().cpu().numpy()
# Calculate the norm
uap_norm = np.linalg.norm(uap_numpy.reshape(-1), np.inf)
print_log("Norm of UAP: {}".format(uap_norm), log)
log.close()
def metrics_evaluate(data_loader, target_model, perturbed_model, targeted, target_class, log=None, use_cuda=True):
# switch to evaluate mode
target_model.eval()
perturbed_model.eval()
perturbed_model.module.generator.eval()
perturbed_model.module.target_model.eval()
clean_acc = AverageMeter()
perturbed_acc = AverageMeter()
attack_success_rate = AverageMeter() # Among the correctly classified samples, the ratio of being different from clean prediction (same as gt)
if targeted:
all_to_target_success_rate = AverageMeter() # The ratio of samples going to the sink classes
all_to_target_success_rate_filtered = AverageMeter()
total_num_samples = 0
num_same_classified = 0
num_diff_classified = 0
for input, gt in data_loader:
if use_cuda:
gt = gt.cuda()
input = input.cuda()
# compute output
with torch.no_grad():
clean_output = target_model(input)
pert_output = perturbed_model(input)
correctly_classified_mask = torch.argmax(clean_output, dim=-1).cpu() == gt.cpu()
cl_acc = accuracy(clean_output.data, gt, topk=(1,))
clean_acc.update(cl_acc[0].item(), input.size(0))
pert_acc = accuracy(pert_output.data, gt, topk=(1,))
perturbed_acc.update(pert_acc[0].item(), input.size(0))
# Calculating Fooling Ratio params
clean_out_class = torch.argmax(clean_output, dim=-1)
pert_out_class = torch.argmax(pert_output, dim=-1)
total_num_samples += len(clean_out_class)
num_same_classified += torch.sum(clean_out_class == pert_out_class).cpu().numpy()
num_diff_classified += torch.sum(~(clean_out_class == pert_out_class)).cpu().numpy()
if torch.sum(correctly_classified_mask)>0:
with torch.no_grad():
pert_output_corr_cl = perturbed_model(input[correctly_classified_mask])
attack_succ_rate = accuracy(pert_output_corr_cl, gt[correctly_classified_mask], topk=(1,))
attack_success_rate.update(attack_succ_rate[0].item(), pert_output_corr_cl.size(0))
# Calculate Absolute Accuracy Drop
aad_source = clean_acc.avg - perturbed_acc.avg
# Calculate Relative Accuracy Drop
if clean_acc.avg !=0:
rad_source = (clean_acc.avg - perturbed_acc.avg)/clean_acc.avg * 100.
else:
rad_source = 0.
# Calculate fooling ratio
fooling_ratio = num_diff_classified/total_num_samples * 100.
if targeted:
# 2. How many of all samples go the sink class (Only relevant for others loader)
target_cl = torch.ones_like(gt) * target_class
all_to_target_succ_rate = accuracy(pert_output, target_cl, topk=(1,))
all_to_target_success_rate.update(all_to_target_succ_rate[0].item(), pert_output.size(0))
# 3. How many of all samples go the sink class, except gt sink class (Only relevant for others loader)
# Filter all idxs which are not belonging to sink class
non_target_class_idxs = [i != target_class for i in gt]
non_target_class_mask = torch.Tensor(non_target_class_idxs)==True
if torch.sum(non_target_class_mask)>0:
gt_non_target_class = gt[non_target_class_mask]
pert_output_non_target_class = pert_output[non_target_class_mask]
target_cl = torch.ones_like(gt_non_target_class) * target_class
all_to_target_succ_rate_filtered = accuracy(pert_output_non_target_class, target_cl, topk=(1,))
all_to_target_success_rate_filtered.update(all_to_target_succ_rate_filtered[0].item(), pert_output_non_target_class.size(0))
if log:
print_log('\n\t#######################', log)
print_log('\tClean model accuracy: {:.3f}'.format(clean_acc.avg), log)
print_log('\tPerturbed model accuracy: {:.3f}'.format(perturbed_acc.avg), log)
print_log('\tAbsolute Accuracy Drop: {:.3f}'.format(aad_source), log)
print_log('\tRelative Accuracy Drop: {:.3f}'.format(rad_source), log)
print_log('\tAttack Success Rate: {:.3f}'.format(100-attack_success_rate.avg), log)
print_log('\tFooling Ratio: {:.3f}'.format(fooling_ratio), log)
if targeted:
print_log('\tAll --> Target Class {} Prec@1 {:.3f}'.format(target_class, all_to_target_success_rate.avg), log)
print_log('\tAll (w/o sink samples) --> Sink {} Prec@1 {:.3f}'.format(target_class, all_to_target_success_rate_filtered.avg), log)
def train(data_loader,
model,
criterion,
optimizer,
epsilon,
num_iterations,
targeted,
target_class,
log,
print_freq=200,
use_cuda=True):
# train function (forward, backward, update)
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to train mode
model.module.generator.train()
model.module.target_model.eval()
end = time.time()
data_iterator = iter(data_loader)
iteration=0
while (iteration<num_iterations):
try:
input, target = next(data_iterator)
except StopIteration:
# StopIteration is thrown if dataset ends
# reinitialize data loader
data_iterator = iter(data_loader)
input, target = next(data_iterator)
if targeted:
target = torch.ones(input.shape[0], dtype=torch.int64) * target_class
# measure data loading time
data_time.update(time.time() - end)
if use_cuda:
target = target.cuda()
input = input.cuda()
# compute output
if model.module._get_name() == "Inception3":
output, aux_output = model(input)
loss1 = criterion(output, target)
loss2 = criterion(aux_output, target)
loss = loss1 + 0.4*loss2
else:
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
if len(target.shape) > 1:
target = torch.argmax(target, dim=-1)
prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(prec1.item(), input.size(0))
top5.update(prec5.item(), input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Projection
model.module.generator.uap.data = torch.clamp(model.module.generator.uap.data, -epsilon, epsilon)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if iteration % print_freq == 0:
print_log(' Iteration: [{:03d}/{:03d}] '
'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) '
'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
'Loss {loss.val:.4f} ({loss.avg:.4f}) '
'Prec@1 {top1.val:.3f} ({top1.avg:.3f}) '
'Prec@5 {top5.val:.3f} ({top5.avg:.3f}) '.format(
iteration, num_iterations, batch_time=batch_time,
data_time=data_time, loss=losses, top1=top1, top5=top5) + time_string(), log)
iteration+=1
print_log(' **Train** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1,
top5=top5,
error1=100-top1.avg), log)
def train_half_half(sources_data_loader, others_data_loader,
model, target_model, criterion, optimizer, epsilon, num_iterations, log,
print_freq=200, use_cuda=True, patch=False):
# train function (forward, backward, update)
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to train mode
model.module.generator.train()
model.module.target_model.eval()
target_model.eval()
end = time.time()
sources_data_iterator = iter(sources_data_loader)
others_data_iterator = iter(others_data_loader)
iteration=0
while (iteration<num_iterations):
try:
sources_input, sources_target = next(sources_data_iterator)
except StopIteration:
# StopIteration is thrown if dataset ends
# reinitialize data loader
sources_data_iterator = iter(sources_data_loader)
sources_input, sources_target = next(sources_data_iterator)
try:
others_input, others_target = next(others_data_iterator)
except StopIteration:
# StopIteration is thrown if dataset ends
# reinitialize data loader
others_data_iterator = iter(others_data_loader)
others_input, others_target = next(others_data_iterator)
# Concat the two batches
input = torch.cat([sources_input, others_input], dim=0)
target = torch.cat([sources_target, others_target], dim=0)
# measure data loading time
data_time.update(time.time() - end)
if use_cuda:
target = target.cuda()
input = input.cuda()
# compute output
output = model(input)
target_model_output = target_model(input)
loss = criterion(output, target_model_output, target)
# measure accuracy and record loss
if len(target.shape) > 1:
target = torch.argmax(target, dim=-1)
prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(prec1.item(), input.size(0))
top5.update(prec5.item(), input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Project to l-infinity ball
if patch:
model.module.generator.uap.data = torch.clamp(model.module.generator.uap.data, 0, epsilon)
else:
model.module.generator.uap.data = torch.clamp(model.module.generator.uap.data, -epsilon, epsilon)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if iteration % print_freq == 0:
print_log(' Iteration: [{:03d}/{:03d}] '
'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) '
'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
'Loss {loss.val:.4f} ({loss.avg:.4f}) '
'Prec@1 {top1.val:.3f} ({top1.avg:.3f}) '
'Prec@5 {top5.val:.3f} ({top5.avg:.3f}) '.format(
iteration, num_iterations, batch_time=batch_time,
data_time=data_time, loss=losses, top1=top1, top5=top5) + time_string(), log)
iteration+=1
print_log(' **Train** Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f} Error@1 {error1:.3f}'.format(top1=top1,
top5=top5,
error1=100-top1.avg), log)
def metrics_evaluate(source_loader, others_loader, target_model, perturbed_model, source_classes, sink_classes, log=None, use_cuda=True):
if len(sink_classes)!=0:
assert len(source_classes) == len(sink_classes)
# switch to evaluate mode
target_model.eval()
perturbed_model.eval()
for loader, loader_name in zip([source_loader, others_loader],["Source", "Others"]):
clean_acc = AverageMeter()
perturbed_acc = AverageMeter()
attack_success_rate = AverageMeter() # Among the correctly classified samples, the ratio of being different from clean prediction (same as gt)
if len(sink_classes)!=0:
all_to_sink_success_rate = []
source_to_sink_success_rate = []
all_to_sink_success_rate_filtered = []
for i in range(len(sink_classes)):
all_to_sink_success_rate.append(AverageMeter()) # The ratio of samples going to the sink classes
source_to_sink_success_rate.append(AverageMeter())
all_to_sink_success_rate_filtered.append(AverageMeter())
total_num_samples = 0
num_same_classified = 0
num_diff_classified = 0
if len(loader)>0: # For UAP, all classes will be attacked, so others_loader is empty
for input, gt in loader:
if use_cuda:
gt = gt.cuda()
input = input.cuda()
# compute output
with torch.no_grad():
clean_output = target_model(input)
pert_output = perturbed_model(input)
correctly_classified_mask = torch.argmax(clean_output, dim=-1).cpu() == gt.cpu()
cl_acc = accuracy(clean_output.data, gt, topk=(1,))
clean_acc.update(cl_acc[0].item(), input.size(0))
pert_acc = accuracy(pert_output.data, gt, topk=(1,))
perturbed_acc.update(pert_acc[0].item(), input.size(0))
# Calculating Fooling Ratio params
clean_out_class = torch.argmax(clean_output, dim=-1)
pert_out_class = torch.argmax(pert_output, dim=-1)
total_num_samples += len(clean_out_class)
num_same_classified += torch.sum(clean_out_class == pert_out_class).cpu().numpy()
num_diff_classified += torch.sum(~(clean_out_class == pert_out_class)).cpu().numpy()
if torch.sum(correctly_classified_mask)>0:
with torch.no_grad():
pert_output_corr_cl = perturbed_model(input[correctly_classified_mask])
attack_succ_rate = accuracy(pert_output_corr_cl, gt[correctly_classified_mask], topk=(1,))
attack_success_rate.update(attack_succ_rate[0].item(), pert_output_corr_cl.size(0))
# Collect samples from Source go to sink
if len(sink_classes)!=0:
# Iterate over source class and sink class pairs
for cl_idx, (source_cl, sink_cl) in enumerate(zip(source_classes, sink_classes)):
# 1. Check how many of the paired source class got to the sink class (Only relevant for source loader)
# Filter all idxs which belong to the source class
source_cl_idxs = [i == source_cl for i in gt]
source_cl_mask = torch.Tensor(source_cl_idxs)==True
if torch.sum(source_cl_mask)>0:
gt_source_cl = gt[source_cl_mask]
pert_output_source_cl = pert_output[source_cl_mask]
# Create desired target value
target_sink = torch.ones_like(gt_source_cl) * sink_cl
source_to_sink_succ_rate = accuracy(pert_output_source_cl, target_sink, topk=(1,))
source_to_sink_success_rate[cl_idx].update(source_to_sink_succ_rate[0].item(), pert_output_source_cl.size(0))
# 2. How many of all samples go the sink class (Only relevant for others loader)
target_sink = torch.ones_like(gt) * sink_cl
all_to_sink_succ_rate = accuracy(pert_output, target_sink, topk=(1,))
all_to_sink_success_rate[cl_idx].update(all_to_sink_succ_rate[0].item(), pert_output.size(0))
# 3. How many of all samples go the sink class, except gt sink class (Only relevant for others loader)
# Filter all idxs which are not belonging to sink class
non_sink_class_idxs = [i != sink_cl for i in gt]
non_sink_class_mask = torch.Tensor(non_sink_class_idxs)==True
if torch.sum(non_sink_class_mask)>0:
gt_non_sink_class = gt[non_sink_class_mask]
pert_output_non_sink_class = pert_output[non_sink_class_mask]
target_sink = torch.ones_like(gt_non_sink_class) * sink_cl
all_to_sink_succ_rate_filtered = accuracy(pert_output_non_sink_class, target_sink, topk=(1,))
all_to_sink_success_rate_filtered[cl_idx].update(all_to_sink_succ_rate_filtered[0].item(), pert_output_non_sink_class.size(0))
if log:
print_log('\n\t########## {} #############'.format(loader_name), log)
if len(sink_classes)!=0:
for cl_idx, (source_cl, sink_cl) in enumerate(zip(source_classes, sink_classes)):
print_log('\n\tSource {} --> Sink {} Prec@1 {:.3f}'.format(source_cl, sink_cl, source_to_sink_success_rate[cl_idx].avg), log)
print_log('\tAll --> Sink {} Prec@1 {:.3f}'.format(sink_cl, all_to_sink_success_rate[cl_idx].avg), log)
# Average fooling ratio of the non-source classes into the target label
print_log('\tAll (w/o sink samples) --> Sink {} Prec@1 {:.3f}'.format(sink_cl, all_to_sink_success_rate_filtered[cl_idx].avg), log)
else:
print('\n\t########## {} #############'.format(loader_name))
if len(sink_classes)!=0:
for cl_idx, (source_cl, sink_cl) in enumerate(zip(source_classes, sink_classes)):
print('\n\tSource {} --> Sink {} Prec@1 {:.3f}'.format(source_cl, sink_cl, source_to_sink_success_rate[cl_idx].avg))
print('\tAll --> Sink {} Prec@1 {:.3f}'.format(sink_cl, all_to_sink_success_rate[cl_idx].avg))
# Average fooling ratio of the non-source classes into the target label
print('\tAll (w/o sink samples) --> Sink {} Prec@1 {:.3f}'.format(sink_cl, all_to_sink_success_rate_filtered[cl_idx].avg))
def main():
args = parse_arguments()
random.seed(args.pretrained_seed)
torch.manual_seed(args.pretrained_seed)
if args.use_cuda:
torch.cuda.manual_seed_all(args.pretrained_seed)
cudnn.benchmark = True
# get a path for saving the model to be trained
model_path = get_model_path(dataset_name=args.pretrained_dataset,
network_arch=args.pretrained_arch,
random_seed=args.pretrained_seed)
# Init logger
log_file_name = os.path.join(model_path, 'log_seed_{}.txt'.format(args.pretrained_seed))
print("Log file: {}".format(log_file_name))
log = open(log_file_name, 'w')
print_log('save path : {}'.format(model_path), log)
state = {k: v for k, v in args._get_kwargs()}
for key, value in state.items():
print_log("{} : {}".format(key, value), log)
print_log("Random Seed: {}".format(args.pretrained_seed), log)
print_log("Python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("Torch version : {}".format(torch.__version__), log)
print_log("Cudnn version : {}".format(torch.backends.cudnn.version()), log)
# Get data specs
num_classes, (mean, std), input_size, num_channels = get_data_specs(args.pretrained_dataset, args.pretrained_arch)
pretrained_data_train, pretrained_data_test = get_data(args.pretrained_dataset,
mean=mean,
std=std,
input_size=input_size,
train_target_model=True)
pretrained_data_train_loader = torch.utils.data.DataLoader(pretrained_data_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
pretrained_data_test_loader = torch.utils.data.DataLoader(pretrained_data_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print_log("=> Creating model '{}'".format(args.pretrained_arch), log)
# Init model, criterion, and optimizer
net = get_network(args.pretrained_arch, input_size=input_size, num_classes=num_classes, finetune=args.finetune)
print_log("=> Network :\n {}".format(net), log)
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
non_trainale_params = get_num_non_trainable_parameters(net)
trainale_params = get_num_trainable_parameters(net)
total_params = get_num_parameters(net)
print_log("Trainable parameters: {}".format(trainale_params), log)
print_log("Non Trainable parameters: {}".format(non_trainale_params), log)
print_log("Total # parameters: {}".format(total_params), log)
# define loss function (criterion) and optimizer
criterion_xent = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), state['learning_rate'], momentum=state['momentum'],
weight_decay=state['decay'], nesterov=True)
if args.use_cuda:
net.cuda()
criterion_xent.cuda()
recorder = RecorderMeter(args.epochs)
# Main loop
start_time = time.time()
epoch_time = AverageMeter()
for epoch in range(args.epochs):
current_learning_rate = adjust_learning_rate(args.learning_rate, args.momentum, optimizer, epoch, args.gammas, args.schedule)
need_hour, need_mins, need_secs = convert_secs2time(epoch_time.avg * (args.epochs-epoch))
need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(need_hour, need_mins, need_secs)
print_log('\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [learning_rate={:6.4f}]'.format(time_string(), epoch, args.epochs, need_time, current_learning_rate) \
+ ' [Best : Accuracy={:.2f}, Error={:.2f}]'.format(recorder.max_accuracy(False), 100-recorder.max_accuracy(False)), log)
# train for one epoch
train_acc, train_los = train_target_model(pretrained_data_train_loader, net, criterion_xent, optimizer, epoch, log,
print_freq=args.print_freq,
use_cuda=args.use_cuda)
# evaluate on validation set
print_log("Validation on pretrained test dataset:", log)
val_acc = validate(pretrained_data_test_loader, net, criterion_xent, log, use_cuda=args.use_cuda)
is_best = recorder.update(epoch, train_los, train_acc, 0., val_acc)
save_checkpoint({
'epoch' : epoch + 1,
'arch' : args.pretrained_arch,
'state_dict' : net.state_dict(),
'recorder' : recorder,
'optimizer' : optimizer.state_dict(),
'args' : copy.deepcopy(args),
}, model_path, 'checkpoint.pth.tar')
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
recorder.plot_curve(os.path.join(model_path, 'curve.png') )
log.close()
def train():
train_data = ACNet_data.SUNRGBD(transform=transforms.Compose([ACNet_data.scaleNorm(),
ACNet_data.RandomScale((1.0, 1.4)),
ACNet_data.RandomHSV((0.9, 1.1),
(0.9, 1.1),
(25, 25)),
ACNet_data.RandomCrop(image_h, image_w),
ACNet_data.RandomFlip(),
ACNet_data.ToTensor(),
ACNet_data.Normalize()]),
phase_train=True,
data_dir=args.data_dir)
train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=False)
num_train = len(train_data)
if args.last_ckpt:
model = ACNet_models_V1.ACNet(num_class=40, pretrained=False)
else:
model = ACNet_models_V1.ACNet(num_class=40, pretrained=True)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
CEL_weighted = utils.CrossEntropyLoss2d(weight=nyuv2_frq)
model.train()
model.to(device)
CEL_weighted.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
momentum=args.momentum, weight_decay=args.weight_decay)
global_step = 0
if args.last_ckpt:
global_step, args.start_epoch = load_ckpt(model, optimizer, args.last_ckpt, device)
lr_decay_lambda = lambda epoch: args.lr_decay_rate ** (epoch // args.lr_epoch_per_decay)
scheduler = LambdaLR(optimizer, lr_lambda=lr_decay_lambda)
writer = SummaryWriter(args.summary_dir)
for epoch in range(int(args.start_epoch), args.epochs):
scheduler.step(epoch)
local_count = 0
last_count = 0
end_time = time.time()
if epoch % args.save_epoch_freq == 0 and epoch != args.start_epoch:
save_ckpt(args.ckpt_dir, model, optimizer, global_step, epoch,
local_count, num_train)
for batch_idx, sample in enumerate(train_loader):
image = sample['image'].to(device)
depth = sample['depth'].to(device)
target_scales = [sample[s].to(device) for s in ['label', 'label2', 'label3', 'label4', 'label5']]
optimizer.zero_grad()
pred_scales = model(image, depth, args.checkpoint)
loss = CEL_weighted(pred_scales, target_scales)
loss.backward()
optimizer.step()
local_count += image.data.shape[0]
global_step += 1
if global_step % args.print_freq == 0 or global_step == 1:
time_inter = time.time() - end_time
count_inter = local_count - last_count
print_log(global_step, epoch, local_count, count_inter,
num_train, loss, time_inter)
end_time = time.time()
for name, param in model.named_parameters():
writer.add_histogram(name, param.clone().cpu().data.numpy(), global_step, bins='doane')
grid_image = make_grid(image[:3].clone().cpu().data, 3, normalize=True)
writer.add_image('image', grid_image, global_step)
grid_image = make_grid(depth[:3].clone().cpu().data, 3, normalize=True)
writer.add_image('depth', grid_image, global_step)
grid_image = make_grid(utils.color_label(torch.max(pred_scales[0][:3], 1)[1] + 1), 3, normalize=False,
range=(0, 255))
writer.add_image('Predicted label', grid_image, global_step)
grid_image = make_grid(utils.color_label(target_scales[0][:3]), 3, normalize=False, range=(0, 255))
writer.add_image('Groundtruth label', grid_image, global_step)
writer.add_scalar('CrossEntropyLoss', loss.data, global_step=global_step)
writer.add_scalar('Learning rate', scheduler.get_lr()[0], global_step=global_step)
last_count = local_count
save_ckpt(args.ckpt_dir, model, optimizer, global_step, args.epochs,
0, num_train)
print("Training completed ")
def main():
args = parse_arguments()
random.seed(args.seed)
torch.manual_seed(args.seed)
if args.use_cuda:
torch.cuda.manual_seed_all(args.seed)
cudnn.benchmark = True
model_path = get_model_path(args.dataset, args.arch, args.seed)
# Init logger
log_file_name = os.path.join(model_path, 'log.txt')
print("Log file: {}".format(log_file_name))
log = open(log_file_name, 'w')
print_log('model path : {}'.format(model_path), log)
state = {k: v for k, v in args._get_kwargs()}
for key, value in state.items():
print_log("{} : {}".format(key, value), log)
print_log("Random Seed: {}".format(args.seed), log)
print_log("Python version : {}".format(sys.version.replace('\n', ' ')),
log)
print_log("Torch version : {}".format(torch.__version__), log)
print_log("Cudnn version : {}".format(torch.backends.cudnn.version()),
log)
# Data specifications for the webistes dataset
mean = [0., 0., 0.]
std = [1., 1., 1.]
input_size = 224
num_classes = 4
# Dataset
traindir = os.path.join(WEBSITES_DATASET_PATH, 'train')
valdir = os.path.join(WEBSITES_DATASET_PATH, 'val')
train_transform = transforms.Compose([
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
test_transform = transforms.Compose([
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
data_train = dset.ImageFolder(root=traindir, transform=train_transform)
data_test = dset.ImageFolder(root=valdir, transform=test_transform)
# Dataloader
data_train_loader = torch.utils.data.DataLoader(data_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
data_test_loader = torch.utils.data.DataLoader(data_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# Network
if args.arch == "vgg16":
net = models.vgg16(pretrained=True)
elif args.arch == "vgg19":
net = models.vgg19(pretrained=True)
elif args.arch == "resnet18":
net = models.resnet18(pretrained=True)
elif args.arch == "resnet50":
net = models.resnet50(pretrained=True)
elif args.arch == "resnet101":
net = models.resnet101(pretrained=True)
elif args.arch == "resnet152":
net = models.resnet152(pretrained=True)
else:
raise ValueError("Network {} not supported".format(args.arch))
if num_classes != 1000:
net = manipulate_net_architecture(model_arch=args.arch,
net=net,
num_classes=num_classes)
# Loss function
if args.loss_function == "ce":
criterion = torch.nn.CrossEntropyLoss()
else:
raise ValueError
# Cuda
if args.use_cuda:
net.cuda()
criterion.cuda()
# Optimizer
momentum = 0.9
decay = 5e-4
optimizer = torch.optim.SGD(net.parameters(),
lr=args.learning_rate,
momentum=momentum,
weight_decay=decay,
nesterov=True)
recorder = RecorderMeter(args.epochs)
start_time = time.time()
epoch_time = AverageMeter()
# Main loop
for epoch in range(args.epochs):
current_learning_rate = adjust_learning_rate(args.learning_rate,
momentum, optimizer,
epoch, args.gammas,
args.schedule)
need_hour, need_mins, need_secs = convert_secs2time(
epoch_time.avg * (args.epochs - epoch))
need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(
need_hour, need_mins, need_secs)
print_log('\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [learning_rate={:6.4f}]'.format(time_string(), epoch, args.epochs, need_time, current_learning_rate) \
+ ' [Best : Accuracy={:.2f}, Error={:.2f}]'.format(recorder.max_accuracy(False), 100-recorder.max_accuracy(False)), log)
# train for one epoch
train_acc, train_los = train_model(data_loader=data_train_loader,
model=net,
criterion=criterion,
optimizer=optimizer,
epoch=epoch,
log=log,
print_freq=200,
use_cuda=True)
# evaluate on test set
print_log("Validation on test dataset:", log)
val_acc, val_loss = validate(data_test_loader,
net,
criterion,
log=log,
use_cuda=args.use_cuda)
recorder.update(epoch, train_los, train_acc, val_loss, val_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict(),
'args': copy.deepcopy(args),
}, model_path, 'checkpoint.pth.tar')
# measure elapsed time
epoch_time.update(time.time() - start_time)
start_time = time.time()
recorder.plot_curve(os.path.join(model_path, 'curve.png'))
log.close()
# Average PSNR on one epoch train_data
train_psnr = sum(psnr_list) / len(psnr_list)
train_loss = sum(running_loss) / len(running_loss)
print('MSE loss: {0:.4f}, SSIM loss: {1:.4f}, SSIM linear loss: {2:.4f}'.
format(np.mean(loss_mse_list), np.mean(loss_ssim_list),
np.mean(loss_ssim_lin_list)))
# save the network parameters
torch.save(
net.state_dict(),
'./weight_checkpoint/rawvsr_{}_{}_{}_{}_{}_raw.pkl'.format(
opt['train']['scale'], opt['network']['nframes'],
opt['network']['nf'], opt['network']['back_RBs'], epoch))
# use evaluation models during the net evaluating
net.eval()
test_psnr = test(net, test_data_loader, device, dataset_name)
one_epoch_time = time.time() - start_time
print_log(epoch + 1, num_epochs, one_epoch_time, train_psnr, test_psnr,
dataset_name)
if test_psnr >= previous_test_psnr:
torch.save(
net.state_dict(),
'./weight_checkpoint/rawvsr_{}_{}_{}_{}_raw_best.pkl'.format(
opt['train']['scale'], opt['network']['nframes'],
opt['network']['nf'], opt['network']['back_RBs']))
previous_test_psnr = test_psnr
