def test_vgg_variant():
x = torch.randn(16, 3, 32, 32).cuda()
model = vgg11_variant(32, 0.3).cuda()
out = model(x)
assert out.size(0) == 16
assert out.size(1) == 512
print(out.shape)
model = vgg13_variant(32, 0.3).cuda()
out = model(x)
assert out.size(0) == 16
assert out.size(1) == 512
print(out.shape)
model = vgg16_variant(32, 0.3).cuda()
out = model(x)
assert out.size(0) == 16
assert out.size(1) == 512
print(out.shape)
model = vgg19_variant(32, 0.3).cuda()
out = model(x)
assert out.size(0) == 16
assert out.size(1) == 512
print(out.shape)
def test_model():
x = torch.randn(16, 3, 32, 32).cuda()
features = vgg11_variant(32, 0.3).cuda()
model = SelectiveNet(features, 512, 10).cuda()
out_class, out_select, out_aux = model(x)
assert out_class.size(0) == 16
assert out_class.size(1) == 10
assert out_select.size(0) == 16
assert out_select.size(1) == 1
assert out_aux.size(0) == 16
assert out_aux.size(1) == 10
print('out_class', out_class.shape)
print('out_select', out_select.shape)
print('out_aux', out_aux.shape)
features = vgg13_variant(32, 0.3).cuda()
model = SelectiveNet(features, 512, 10).cuda()
out_class, out_select, out_aux = model(x)
assert out_class.size(0) == 16
assert out_class.size(1) == 10
assert out_select.size(0) == 16
assert out_select.size(1) == 1
assert out_aux.size(0) == 16
assert out_aux.size(1) == 10
print('out_class', out_class.shape)
print('out_select', out_select.shape)
print('out_aux', out_aux.shape)
features = vgg16_variant(32, 0.3).cuda()
model = SelectiveNet(features, 512, 10).cuda()
out_class, out_select, out_aux = model(x)
assert out_class.size(0) == 16
assert out_class.size(1) == 10
assert out_select.size(0) == 16
assert out_select.size(1) == 1
assert out_aux.size(0) == 16
assert out_aux.size(1) == 10
print('out_class', out_class.shape)
print('out_select', out_select.shape)
print('out_aux', out_aux.shape)
features = vgg19_variant(32, 0.3).cuda()
model = SelectiveNet(features, 512, 10).cuda()
out_class, out_select, out_aux = model(x)
assert out_class.size(0) == 16
assert out_class.size(1) == 10
assert out_select.size(0) == 16
assert out_select.size(1) == 1
assert out_aux.size(0) == 16
assert out_aux.size(1) == 10
print('out_class', out_class.shape)
print('out_select', out_select.shape)
print('out_aux', out_aux.shape)
def test_selective_loss():
x = torch.randn(16, 3, 32, 32).cuda()
features = vgg16_variant(32, 0.3).cuda()
model = SelectiveNet(features, 512, 10).cuda()
out_class, out_select, _ = model(x)
target = torch.randint(0, 9, (16, )).cuda()
loss_func = torch.nn.CrossEntropyLoss(reduction='none')
loss, loss_dict = SelectiveLoss(loss_func,
coverage=0.7)(out_class, out_select,
target)
assert loss.size(0) == 1
print('loss', loss)
print(loss_dict)
prediction_out = self.classifier(x)
selection_out = self.selector(x)
auxiliary_out = self.aux_classifier(x)
return prediction_out, selection_out, auxiliary_out
def _initialize_weights(self, module):
for m in module.modules():
if isinstance(m, torch.nn.Conv2d):
torch.nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None:
torch.nn.init.constant_(m.bias, 0)
elif isinstance(m, torch.nn.BatchNorm1d):
torch.nn.init.constant_(m.weight, 1)
torch.nn.init.constant_(m.bias, 0)
elif isinstance(m, torch.nn.Linear):
torch.nn.init.normal_(m.weight, 0, 0.01)
torch.nn.init.constant_(m.bias, 0)
if __name__ == '__main__':
import os
import sys
base = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../')
sys.path.append(base)
from selectivenet.vgg_variant import vgg16_variant
features = vgg16_variant(32,0.3).cuda()
model = SelectiveNet(features,512,10).cuda()
from selectivenet.loss import SelectiveLoss
from selectivenet.data import DatasetBuilder
from external.dada.io import load_model
# dataset
dataset_builder = DatasetBuilder(name='cifar10',
root_path='../../../../data')
test_dataset = dataset_builder(train=False, normalize=True)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=8,
shuffle=False,
num_workers=16,
pin_memory=True)
# model
features = vgg16_variant(dataset_builder.input_size, 0.3).cuda()
model = SelectiveNet(features, 512, dataset_builder.num_classes).cuda()
weight_path = '../../../../logs/abci/cifar10_ex01/0f7d6935b74b42c0a705dd241635d76b/weight_final_coverage_0.70.pth'
load_model(model, weight_path)
if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model)
# attacker
attacker = PGDAttackVariant(10,
32,
30,
'cifar10',
coverage=0.7,
trg_loss='both')
# test
def train(**kwargs):
FLAGS = FlagHolder()
FLAGS.initialize(**kwargs)
FLAGS.summary()
FLAGS.dump(
path=os.path.join(FLAGS.log_dir, 'flags{}.json'.format(FLAGS.suffix)))
# dataset
dataset_builder = DatasetBuilder(name=FLAGS.dataset,
root_path=FLAGS.dataroot)
train_dataset = dataset_builder(train=True, normalize=FLAGS.normalize)
val_dataset = dataset_builder(train=False, normalize=FLAGS.normalize)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=FLAGS.batch_size,
shuffle=True,
num_workers=FLAGS.num_workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=FLAGS.batch_size,
shuffle=False,
num_workers=FLAGS.num_workers,
pin_memory=True)
# model
features = vgg16_variant(dataset_builder.input_size,
FLAGS.dropout_prob).cuda()
model = SelectiveNet(features, FLAGS.dim_features,
dataset_builder.num_classes).cuda()
if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model)
# optimizer
params = model.parameters()
optimizer = torch.optim.SGD(params,
lr=FLAGS.lr,
momentum=FLAGS.momentum,
weight_decay=FLAGS.wd)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=25,
gamma=0.5)
# loss
base_loss = torch.nn.CrossEntropyLoss(reduction='none')
SelectiveCELoss = SelectiveLoss(base_loss, coverage=FLAGS.coverage)
# logger
train_logger = Logger(path=os.path.join(
FLAGS.log_dir, 'train_log{}.csv'.format(FLAGS.suffix)),
mode='train')
val_logger = Logger(path=os.path.join(
FLAGS.log_dir, 'val_log{}.csv'.format(FLAGS.suffix)),
mode='val')
for ep in range(FLAGS.num_epochs):
# pre epoch
train_metric_dict = MetricDict()
val_metric_dict = MetricDict()
# train
for i, (x, t) in enumerate(train_loader):
model.train()
x = x.to('cuda', non_blocking=True)
t = t.to('cuda', non_blocking=True)
# forward
out_class, out_select, out_aux = model(x)
# compute selective loss
loss_dict = OrderedDict()
# loss dict includes, 'empirical_risk' / 'emprical_coverage' / 'penulty'
selective_loss, loss_dict = SelectiveCELoss(
out_class, out_select, t)
selective_loss *= FLAGS.alpha
loss_dict['selective_loss'] = selective_loss.detach().cpu().item()
# compute standard cross entropy loss
ce_loss = torch.nn.CrossEntropyLoss()(out_aux, t)
ce_loss *= (1.0 - FLAGS.alpha)
loss_dict['ce_loss'] = ce_loss.detach().cpu().item()
# total loss
loss = selective_loss + ce_loss
loss_dict['loss'] = loss.detach().cpu().item()
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_metric_dict.update(loss_dict)
# validation
with torch.autograd.no_grad():
for i, (x, t) in enumerate(val_loader):
model.eval()
x = x.to('cuda', non_blocking=True)
t = t.to('cuda', non_blocking=True)
# forward
out_class, out_select, out_aux = model(x)
# compute selective loss
loss_dict = OrderedDict()
# loss dict includes, 'empirical_risk' / 'emprical_coverage' / 'penulty'
selective_loss, loss_dict = SelectiveCELoss(
out_class, out_select, t)
selective_loss *= FLAGS.alpha
loss_dict['selective_loss'] = selective_loss.detach().cpu(
).item()
# compute standard cross entropy loss
ce_loss = torch.nn.CrossEntropyLoss()(out_aux, t)
ce_loss *= (1.0 - FLAGS.alpha)
loss_dict['ce_loss'] = ce_loss.detach().cpu().item()
# total loss
loss = selective_loss + ce_loss
loss_dict['loss'] = loss.detach().cpu().item()
# evaluation
evaluator = Evaluator(out_class.detach(), t.detach(),
out_select.detach())
loss_dict.update(evaluator())
val_metric_dict.update(loss_dict)
# post epoch
# print_metric_dict(ep, FLAGS.num_epochs, train_metric_dict.avg, mode='train')
print_metric_dict(ep,
FLAGS.num_epochs,
val_metric_dict.avg,
mode='val')
train_logger.log(train_metric_dict.avg, step=(ep + 1))
val_logger.log(val_metric_dict.avg, step=(ep + 1))
scheduler.step()
# post training
save_model(model,
path=os.path.join(FLAGS.log_dir,
'weight_final{}.pth'.format(FLAGS.suffix)))
def test(**kwargs):
FLAGS = FlagHolder()
FLAGS.initialize(**kwargs)
FLAGS.summary()
# dataset
dataset_builder = DatasetBuilder(name=FLAGS.dataset,
root_path=FLAGS.dataroot)
test_dataset = dataset_builder(train=False, normalize=FLAGS.normalize)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=FLAGS.batch_size,
shuffle=False,
num_workers=FLAGS.num_workers,
pin_memory=True)
# model
features = vgg16_variant(dataset_builder.input_size,
FLAGS.dropout_prob).cuda()
model = SelectiveNet(features, FLAGS.dim_features,
dataset_builder.num_classes).cuda()
load_model(model, FLAGS.weight)
if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model)
# loss
base_loss = torch.nn.CrossEntropyLoss(reduction='none')
SelectiveCELoss = SelectiveLoss(base_loss, coverage=FLAGS.coverage)
# pre epoch
test_metric_dict = MetricDict()
# test
with torch.autograd.no_grad():
for i, (x, t) in enumerate(test_loader):
model.eval()
x = x.to('cuda', non_blocking=True)
t = t.to('cuda', non_blocking=True)
# forward
out_class, out_select, out_aux = model(x)
# compute selective loss
loss_dict = OrderedDict()
# loss dict includes, 'empirical_risk' / 'emprical_coverage' / 'penulty'
selective_loss, loss_dict = SelectiveCELoss(
out_class, out_select, t)
selective_loss *= FLAGS.alpha
loss_dict['selective_loss'] = selective_loss.detach().cpu().item()
# compute standard cross entropy loss
ce_loss = torch.nn.CrossEntropyLoss()(out_aux, t)
ce_loss *= (1.0 - FLAGS.alpha)
loss_dict['ce_loss'] = ce_loss.detach().cpu().item()
# total loss
loss = selective_loss + ce_loss
loss_dict['loss'] = loss.detach().cpu().item()
# evaluation
evaluator = Evaluator(out_class.detach(), t.detach(),
out_select.detach())
loss_dict.update(evaluator())
test_metric_dict.update(loss_dict)
# post epoch
print_metric_dict(None, None, test_metric_dict.avg, mode='test')
return test_metric_dict.avg