def __init__(self, dropout_rate, feat_length=512, archi_type='resnet18'):
super(CIFAR10FeatureLayer, self).__init__()
self.archi_type = archi_type
self.feat_length = feat_length
if self.archi_type == 'default':
self.add_module('conv1', nn.Conv2d(3, 32, kernel_size=3,
padding=1))
self.add_module('bn1', nn.BatchNorm2d(32))
self.add_module('relu1', nn.ReLU())
self.add_module('pool1', nn.MaxPool2d(kernel_size=2))
#self.add_module('drop1', nn.Dropout(dropout_rate))
self.add_module('conv2', nn.Conv2d(32,
32,
kernel_size=3,
padding=1))
self.add_module('bn2', nn.BatchNorm2d(32))
self.add_module('relu2', nn.ReLU())
self.add_module('pool2', nn.MaxPool2d(kernel_size=2))
#self.add_module('drop2', nn.Dropout(dropout_rate))
self.add_module('conv3', nn.Conv2d(32,
64,
kernel_size=3,
padding=1))
self.add_module('bn3', nn.BatchNorm2d(64))
self.add_module('relu3', nn.ReLU())
self.add_module('pool3', nn.MaxPool2d(kernel_size=2))
#self.add_module('drop3', nn.Dropout(dropout_rate))
elif self.archi_type == 'resnet18':
self.add_module('resnet18', resnet.ResNet18(feat_length))
elif self.archi_type == 'resnet50':
self.add_module('resnet50', resnet.ResNet50(feat_length))
elif self.archi_type == 'resnet152':
self.add_module('resnet152', resnet.ResNet152(feat_length))
else:
raise NotImplementedError
def __init__(self, num_classes):
super().__init__(num_classes)
self.num_classes = num_classes
self.resnet = resnet.ResNet18(num_classes=512)
self.fc1_drop = mc_dropout.MCDropout()
self.fc1 = nn.Linear(512, 256)
self.fc2_drop = mc_dropout.MCDropout()
self.fc2 = nn.Linear(256, num_classes)
def get_model(args):
if args.model == 'resnet18':
model = resnet.ResNet18(num_classes=10)
elif args.model == 'resnet50':
model = resnet.ResNet50(num_classes=10)
elif args.model == 'densenet40':
model = densenet.DenseNet3(depth=40, num_classes=10)
elif args.model == 'vgg16':
model = vgg.VGG('VGG16')
return model
def __init__(self, dataset='mnist'):
super(resnet18, self).__init__()
self.input_height = 28
self.input_width = 28
self.input_dim = 1
self.class_num = 10
self.input_shape = (self.input_dim, self.input_height,
self.input_width)
self.model = resnet.ResNet18(self.input_shape, self.class_num)
def __init__(self,
in_ch,
out_ch,
epochs,
device,
log_interval,
start_lr=2e-4,
lmbda=10.0,
idt_coef=0.5,
g_ch=64,
d_ch=64,
decay_epoch=0):
self.eval = resnet.ResNet18()
self.pool = nn.AvgPool2d(4, 4)
self.eval.load_state_dict(
torch.load('./parameter.pkl', map_location='cpu'))
self.epochs = epochs
self.decay_epoch = decay_epoch if decay_epoch > 0 else int(
self.epochs / 2)
self.lmbda = lmbda
self.log_interval = log_interval
self.idt_coef = idt_coef
self.device = device
self.netG_S2P = Generator(in_ch, out_ch, g_ch)
self.netG_P2S = Generator(in_ch, out_ch, g_ch)
self.netD_S = Discriminator(in_ch, d_ch)
self.netD_P = Discriminator(in_ch, d_ch)
self.init_models()
self.mse_loss = nn.MSELoss()
self.l1_loss = nn.L1Loss()
self.adam_gen = torch.optim.Adam(itertools.chain(
self.netG_S2P.parameters(), self.netG_P2S.parameters()),
lr=start_lr,
betas=(0.5, 0.999))
self.adam_des = torch.optim.Adam(itertools.chain(
self.netD_S.parameters(), self.netD_P.parameters()),
lr=start_lr,
betas=(0.5, 0.999))
self.sample_style = sample_fake()
self.sample_photo = sample_fake()
gen_lr = lr_sched(self.decay_epoch, self.epochs)
des_lr = lr_sched(self.decay_epoch, self.epochs)
self.gen_lr_sched = torch.optim.lr_scheduler.LambdaLR(
self.adam_gen, gen_lr.step)
self.des_lr_sched = torch.optim.lr_scheduler.LambdaLR(
self.adam_des, des_lr.step)
def __init__(self,
n_classes,
init_lr=0.1,
momentum=0.9,
weight_decay=5e-4,
device='cuda',
log_dir='',
ckpt_file='',
model='resnet-10',
multi_gpu=True):
super().__init__(n_classes, init_lr, momentum, weight_decay, device)
self.n_planes = [64, 128, 256, 512]
if model == 'resnet-10':
self.net = resnet.ResNet10(n_classes=self.n_classes,
n_output_planes=self.n_planes)
if model == 'resnet-18':
self.net = resnet.ResNet18(n_classes=self.n_classes,
n_output_planes=self.n_planes)
elif model == 'resnet-34':
self.net = resnet.ResNet34(n_classes=self.n_classes,
n_output_planes=self.n_planes)
elif model == 'resnet-50':
self.net = resnet.ResNet50(n_classes=self.n_classes,
n_output_planes=self.n_planes)
self.net.to(self.device)
if ckpt_file:
print('loading pretrained classifier checkpoint')
if device == 'cpu':
ckpt = torch.load(ckpt_file,
map_location=lambda storage, loc: storage)
else:
ckpt = torch.load(ckpt_file)
self.net.load_state_dict(ckpt['clf'])
if multi_gpu and self.device == 'cuda':
print('replicating model on multiple gpus ... ')
self.net = torch.nn.DataParallel(self.net)
self.optim = torch.optim.SGD(self.net.parameters(),
self.init_lr,
momentum=self.momentum,
weight_decay=self.weight_decay)
self.criterion = torch.nn.CrossEntropyLoss().to(self.device)
print('Number of dnn parameters: {}'.format(
sum([p.data.nelement() for p in self.net.parameters()])))
if log_dir:
utils.save_model_desc(self.net, join(log_dir,
'classifier_desc.txt'))
def get_net(name, args, strategy):
if name == 'resnet18':
net = resnet.ResNet18(strategy, args)
elif name == 'resnet34':
net = ResNet34(n_classes=args['num_classes'],
n_channels=args['num_channels'],
device=args['device'])
elif name == 'net3':
net = Net3(args)
elif name == 'net5':
net = COAPModNet(args)
else:
return "Invalid name"
return net
def main():
# Resnet18 - Reference for CIFAR 10
ref_params = 5586981
ref_flops = 834362880
# WideResnet-28-10 - Reference for CIFAR 100
# ref_params = 36500000
# ref_flops = 10490000000
model = resnet.ResNet18()
print(model)
flops, params = profile(model, (1, 3, 32, 32))
flops, params = flops.item(), params.item()
score_flops = flops / ref_flops
score_params = params / ref_params
score = score_flops + score_params
print("Flops: {}, Params: {}".format(flops, params))
print("Score flops: {} Score Params: {}".format(score_flops, score_params))
print("Final score: {}".format(score))
def get_model(model_name, parameters):
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
if model_name == 'resnet18':
net = resnet.ResNet18(parameters, num_classes=10).to(device)
elif model_name == 'resnet34':
net = resnet.ResNet34(parameters, num_classes=10).to(device)
elif model_name == 'resnet50':
net = resnet.ResNet50(parameters, num_classes=10).to(device)
elif model_name == 'resnet101':
net = resnet.ResNet101(parameters, num_classes=10).to(device)
elif model_name == 'resnet152':
net = resnet.ResNet152(parameters, num_classes=10).to(device)
elif model_name == 'vgg16':
net = 0
else:
print("Entered student model is not compatibale currently!\n")
net = -1
return net
def get_model(device):
"""
:param device: instance of torch.device
:return: An instance of torch.nn.Module
"""
num_classes = 2
if config["dataset"] == "Cifar100":
num_classes = 100
elif config["dataset"] == "Cifar10":
num_classes = 10
elif config["dataset"] == "15-Scene":
num_classes = 15
elif config["dataset"] == "MNIST":
num_classes = 10
model = {
"resnet10":
lambda: resnet.ResNet10(num_classes=num_classes),
"resnet18":
lambda: resnet.ResNet18(num_classes=num_classes),
"resnet34":
lambda: resnet.ResNet34(num_classes=num_classes),
"resnet50":
lambda: resnet.ResNet50(num_classes=num_classes),
"resnet101":
lambda: resnet.ResNet101(num_classes=num_classes),
"resnet152":
lambda: resnet.ResNet152(num_classes=num_classes),
"bert":
lambda: modeling_bert_appendix.BertImage(config,
num_classes=num_classes),
}[config["model"]]()
model.to(device)
if device == "cuda":
# model = torch.nn.DataParallel(model) # multiple GPUs not available
# for free on Google Colab -EU
torch.backends.cudnn.benchmark = True
return model
def main(args):
#resnet_blocks = [3,8,36,3]
#resnet_blocks = [3,4,23,3]
#resnet_blocks = [3,4,6,3]
#resnet_blocks = [2,2,2,2]
input_x = np.random.randn(1,3,32,32).astype(np.float32)
#input_x = np.random.randn(1,3,416,416).astype(np.float32)
input_x = torch.autograd.Variable(torch.from_numpy(input_x)).cpu()
tests = [
('simplenet7_thin', simplenet.SimpleNet7_thin()),
('simplenet7', simplenet.SimpleNet7()),
('simplenet9_mobile', simplenet.SimpleNet9_mobile()),
('simplenet9_thin', simplenet.SimpleNet9_thin()),
('simplenet9', simplenet.SimpleNet9()),
('resnet18', resnet.ResNet18()),
('resnet18JPEG', resnet_with_compression.ResNet18()),
('mobilenetv2', mobilenetv2.MobileNetV2()),
('mobilenet', mobilenet.MobileNet()),
]
for test in tests:
net_name = test[0]
net = test[1].eval()
time = []
for _ in range(N):
t1 = timeit.default_timer()
output = net(input_x)
t2 = timeit.default_timer()
time.append(t2-t1)
time = np.asarray(time)
sys.stdout.write(net_name + ',' + str(np.median(time)) + ',' + str(np.average(time)) + ',' + str(np.std(time))+'\n')
return
def main():
args.split_data = False
args.distributed = True
print("~~epoch\thours\ttop1Accuracy\n")
start_time = datetime.now()
if args.distributed:
os.environ['WORLD_SIZE'] = str(args.world_size)
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=int(os.environ['RANK']))
torch.cuda.set_device(args.local_rank)
if dist.get_rank() == 0:
print(str(dist.get_world_size()) + ' number of workers is set up!')
if dist.get_rank() == 0:
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
log_writer = tensorboardX.SummaryWriter(
args.save_dir) if dist.get_rank() == 0 else None
# create model for cifar-10
model = resnet.ResNet18(num_classes=10)
#model = models.vgg11_bn(10)
model = model.cuda()
n_dev = torch.cuda.device_count()
#model.para sync
global param_copy
param_copy = list(model.parameters())
for parameter in param_copy:
dist.broadcast(parameter.data, 0, group=0)
if dist.get_rank() == 0:
print('parameter sync finished')
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
if args.communication_method == 'Signum':
optimizer = Signum_optimizer.SGD_distribute(
param_copy,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
local_rank=args.local_rank,
compression_buffer=args.compress,
all_reduce=args.all_reduce,
args=args)
elif args.communication_method == 'QSGD':
optimizer = QSGD_optimizer.SGD_distribute(
param_copy,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
local_rank=args.local_rank,
compression_buffer=args.compress,
all_reduce=args.all_reduce,
args=args)
'''
if args.enable_max:
optimizer.enable_max_norm()
else:
optimizer.disable_max_norm()
'''
best_prec1 = 0
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
checkpoint = torch.load(
args.resume,
map_location=lambda storage, loc: storage.cuda(args.gpu))
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
print("=> no checkpoint found at '{}'".format(args.resume))
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
args.sz = 224
train_loader, val_loader, train_sampler = get_loaders(
traindir, valdir, split_data=args.split_data, seed=args.seed)
if args.evaluate:
return validate(val_loader, model, criterion, epoch, start_time)
for epoch in range(args.start_epoch, (args.epochs + args.extra_epochs)):
adjust_learning_rate(optimizer, epoch)
if args.distributed:
train_sampler.set_epoch(epoch + args.seed)
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=UserWarning)
train(train_loader, model, criterion, optimizer, epoch, log_writer)
if args.prof: break
prec1 = validate(val_loader, model, criterion, epoch, start_time,
log_writer)
if dist.get_rank() == 0:
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
'''
num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform_test)
testloader = torch.utils.data.DataLoader(testset,
batch_size=100,
shuffle=False,
num_workers=2)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
# Model
print('==> Building model..')
net = resnet.ResNet18()
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
if resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
# address_head = './drive/My Drive/checkpoint/ckpt_resnet18_'
checkpoint = torch.load(address_head + '100.t7')
net.load_state_dict(checkpoint['net'])
best_acc = checkpoint['acc']
start_epoch = checkpoint['epoch']
print('==> Resume Done')
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batchSz', type=int, default=128)
parser.add_argument('--nEpochs', type=int, default=250)
parser.add_argument('--no-cuda', action='store_true')
parser.add_argument('--save')
parser.add_argument('--seed', type=int, default=1)
parser.add_argument('--opt',
type=str,
default='sgd',
choices=('sgd', 'adam', 'rmsprop'))
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
args.save = args.save or 'work/wd/weightdecay0.001'
setproctitle.setproctitle(args.save)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
if os.path.exists(args.save):
shutil.rmtree(args.save)
os.makedirs(args.save, exist_ok=True)
normMean = [0.5423671, 0.53410053, 0.52827841]
normStd = [0.30129549, 0.29579896, 0.29065931]
normTransform = transforms.Normalize(normMean, normStd)
trainTransform = transforms.Compose([
transforms.Scale(64),
transforms.RandomCrop(64, padding=8),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normTransform
])
testTransform = transforms.Compose(
[transforms.Scale(64),
transforms.ToTensor(), normTransform])
kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
trainLoader = DataLoader(dset.CIFAR100(
root='/datadisk/A/liubaoen/data/cifar',
train=True,
download=True,
transform=trainTransform),
batch_size=args.batchSz,
shuffle=True,
**kwargs)
valLoader = DataLoader(dset.CIFAR100(
root='/datadisk/A/liubaoen/data/cifar',
train=False,
download=True,
transform=testTransform),
batch_size=args.batchSz,
shuffle=False,
**kwargs)
net = resnet.ResNet18()
print(' + Number of params: {}'.format(
sum([p.data.nelement() for p in net.parameters()])))
if args.cuda:
net = net.cuda()
net = torch.nn.DataParallel(net,
device_ids=range(
torch.cuda.device_count()))
cudnn.benchmark = True
if args.opt == 'sgd':
optimizer = optim.SGD(net.parameters(),
lr=1e-1,
momentum=0.9,
weight_decay=0.001)
elif args.opt == 'adam':
optimizer = optim.Adam(net.parameters(), weight_decay=1e-4)
elif args.opt == 'rmsprop':
optimizer = optim.RMSprop(net.parameters(), weight_decay=1e-4)
trainF = open(os.path.join(args.save, 'train.csv'), 'w')
testF = open(os.path.join(args.save, 'test.csv'), 'w')
for epoch in range(1, args.nEpochs + 1):
adjust_opt(args.opt, optimizer, epoch)
train(args, epoch, net, trainLoader, optimizer, trainF)
test(args, epoch, net, valLoader, optimizer, testF)
torch.save(net, os.path.join(args.save, 'latest.pth'))
os.system('./plot.py {} &'.format(args.save))
trainF.close()
testF.close()
def __init__(self, args):
# parameters
self.epoch = args.epoch
self.sample_num = 100
self.batch_size = args.batch_size
self.save_dir = args.save_dir
self.result_dir = args.result_dir
self.dataset = args.dataset
self.log_dir = args.log_dir
self.gpu_mode = args.gpu_mode
self.model_name = args.gan_type
self.input_size = args.input_size
self.z_dim = 2352
self.lambda_ = 0.25
self.checkpoint = args.checkpoint
with open(
os.path.join(self.save_dir, self.dataset, self.model_name,
self.model_name + 'best.txt'), 'r') as f:
temp = f.readlines()
self.loss_adv_avg = float(temp[0])
self.loss_perturb_avg = float(temp[1])
# load dataset
self.data_loader = dataloader(self.dataset, self.input_size,
self.batch_size)
data = self.data_loader.__iter__().__next__()[0]
# networks init
self.G = generator(input_dim=self.z_dim,
output_dim=data.shape[1],
input_size=self.input_size)
print(self.z_dim, data.shape[1], self.input_size)
self.D = discriminator(input_dim=data.shape[1],
output_dim=1,
input_size=self.input_size)
save_dir = os.path.join(self.save_dir, self.dataset, self.model_name)
print(save_dir)
# load checkpoint
if self.checkpoint != '':
print(self.checkpoint + 'G.pkl')
self.G.load_state_dict(torch.load(self.checkpoint + 'G.pkl'))
self.D.load_state_dict(torch.load(self.checkpoint + 'D.pkl'))
self.G_optimizer = optim.Adam(self.G.parameters(),
lr=args.lrG,
betas=(args.beta1, args.beta2))
self.D_optimizer = optim.Adam(self.D.parameters(),
lr=args.lrD,
betas=(args.beta1, args.beta2))
if self.gpu_mode:
self.G.cuda()
self.D.cuda()
self.BCE_loss = nn.BCELoss().cuda()
else:
self.BCE_loss = nn.BCELoss()
# targeted model
net = resnet.ResNet18()
net = net.cuda()
net = torch.nn.DataParallel(net)
checkpoint = torch.load("H:/pytorch-cifar/checkpoint/DataPackpt.pth")
net.load_state_dict(checkpoint['net'])
self.model = net
self.model.eval()
print('---------- Networks architecture -------------')
utils.print_network(self.G)
utils.print_network(self.D)
print('-----------------------------------------------')
# fixed noise
self.sample_z_ = torch.rand((self.batch_size, self.z_dim))
if self.gpu_mode:
self.sample_z_ = self.sample_z_.cuda()
import torchvision.transforms as transforms
import os
import argparse
import time
import numpy as np
import resnet
import vgg
import copy
'''
Vectorizing w_70, w_150 and theta
'''
net_A = resnet.ResNet18()
net_B = resnet.ResNet18()
net_C = resnet.ResNet18()
net_A.load_state_dict(torch.load('results/15417101052776022/epoch_70.t7'))
net_A.eval()
net_B.load_state_dict(torch.load('results/15417101052776022/epoch_150.t7'))
net_B.eval()
net_C.load_state_dict(
torch.load('results/curve_find/15418861362468774/theta_step_50499.t7'))
net_C.eval()
# net_C = torch.load('curve_find/15417853945005875/theta_step_50499.t7')
# net_C.eval()
param_vec_A = torch.cat(
import torch
transformer = transforms.Compose([
transforms.Resize((250, 250)),
transforms.RandomHorizontalFlip(),
transforms.RandomAffine(10, (10, 10)),
transforms.RandomResizedCrop((224, 224), scale=(0.85, 1.0)),
transforms.ToTensor()
])
train_dataset = ImageFolder("./data", transform=transformer)
print(train_dataset.class_to_idx)
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=128,
shuffle=True)
model = resnet.ResNet18(class_num=3).cuda()
optimizer = torch.optim.SGD(model.parameters(),
lr=1e-3,
weight_decay=1e-4,
momentum=0.9)
criterion = torch.nn.CrossEntropyLoss().cuda()
model.train()
for epoch in range(30):
for input, target in train_dataloader:
input = torch.Tensor(input).cuda()
target = torch.Tensor(target).long().cuda()
out = model(input)
loss = criterion(out, target)
optimizer.zero_grad()
download=True,
transform=transform)
testset = torchvision.datasets.CIFAR10(root='../data',
train=False,
download=True,
transform=transform)
trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testloader = DataLoader(testset, batch_size=128, shuffle=True, num_workers=2)
dset_loaders = {'train': trainloader, 'val': testloader}
dset_sizes = {
'train': len(trainset.train_labels),
'val': len(testset.test_labels)
}
# Defining the model
model = resnet.ResNet18()
if torch.cuda.is_available():
model.cuda()
criterion = CrossEntropyLoss()
optimizer = Adam(model.parameters(),
lr=0.001,
betas=(0.9, 0.999),
weight_decay=5e-4)
# Training the network
train_outputs = train_model(model,
dset_loaders,
dset_sizes,
criterion,
optimizer,
num_epochs=90,
args = parser.parse_args()
test_data = BPEI_Dataset(args.test_list,
args.data_dir,
loader=BPEI_loader.posemulti_loader)
test_loader = DataLoader(
dataset=test_data,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
drop_last=True,
)
#device_ids = [0, 7, 9]
model = resnet.ResNet18(input_nc=1)
#model = nn.DataParallel(model,device_ids=device_ids).cuda()
model = nn.DataParallel(model).cuda()
model.load_state_dict(torch.load(args.model))
#print(model.state_dict().keys())
#-------weight--------------------------------
# for i, j in model.named_parameters():
# if i.find('gamma')>0:
# print(i)
# print(j)
save_path = args.save_path
if not os.path.exists(save_path):
os.makedirs(save_path)
def main():
# Config
parser = argparse.ArgumentParser()
parser.add_argument("--config", default="./config/example.yaml")
args = parser.parse_args()
assert os.path.exists(args.config)
print("Load configuration file from {}".format(args.config))
with open(args.config, "r") as f:
print(f.read())
f.seek(0)
config = yaml.safe_load(f)
assert os.path.exists(config["saved_path"])
config_name = args.config.split("/")[-1].split(".")[0]
saved_path = os.path.join(config["saved_path"], config_name)
if not os.path.exists(saved_path):
os.mkdir(saved_path)
shutil.copy2(args.config, saved_path)
# Training dataset and dataloader
root = config["dataset_path"]
pratio = config["poison_ratio"]
tlabel = config["target_label"]
trigger_loc = torch.tensor(config["trigger"])
trigger_ptn = torch.randint(0, 256, [len(trigger_loc)])
train_data = cifar.CIFAR10(root)
train_pidx = gen_pidx(train_data.targets, pratio, tlabel)
train_data = cifar.BadCIFAR10(root, train_pidx, tlabel, trigger_loc,
trigger_ptn)
train_data.set_transform(**config["train_transform"])
train_loader = DataLoader(train_data,
batch_size=config["batch_size"],
shuffle=True)
# Normal test dataset and dataloader
test_data = cifar.CIFAR10(root, train=False)
test_data.set_transform(**config["test_transform"])
test_loader = DataLoader(test_data,
batch_size=config["batch_size"],
shuffle=False)
# Backdoor test dataset and dataloader
test_pidx = np.ones(len(test_data.data))
bd_test_data = cifar.BadCIFAR10(root,
test_pidx,
tlabel,
trigger_loc,
trigger_ptn,
train=False)
bd_test_data.set_transform(**config["bd_test_transform"])
bd_test_loader = DataLoader(bd_test_data,
batch_size=config["batch_size"],
shuffle=False)
model = resnet.ResNet18()
model = model.cuda()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"])
log = []
for epoch in range(config["num_epoches"] + 1):
start_time = time.time()
train_results = train(model, train_loader, criterion, optimizer)
bd_results = test(model, bd_test_loader, criterion)
test_results = test(model, test_loader, criterion)
elapse_time = time.time() - start_time
row_dict = {
"epoch": ["{}/{}".format(epoch + 1, config["num_epoches"])],
"train_loss": [train_results["tloss"]],
"train_acc(%)": [train_results["acc"]],
"test_loss": [test_results["loss"]],
"test_acc(%)": [test_results["acc"]],
"bd_loss": [bd_results["loss"]],
"bd_acc": [bd_results["acc"]],
"time(s)": [elapse_time],
}
table = tabulate(row_dict,
headers="keys",
tablefmt="github",
floatfmt="9.5f")
if (epoch + 1) % 30 == 1:
table = table.split("\n")
table = "\n".join([table[1]] + table)
else:
table = table.split("\n")[2]
print(table)
log.append({k: v[0] for k, v in row_dict.items()})
pd.DataFrame(log).to_csv(os.path.join(saved_path, "log.csv"),
index=False)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batchSz', type=int, default=128)
parser.add_argument('--nEpochs', type=int, default=300)
parser.add_argument('--no-cuda', action='store_true')
parser.add_argument('--save')
parser.add_argument('--lr', type=float, default=0.1)
parser.add_argument('--wd', type=float, default=0.0005)
parser.add_argument('--seed', type=int, default=1)
parser.add_argument('--opt',
type=str,
default='sgd',
choices=('sgd', 'adam', 'rmsprop'))
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
args.save = args.save or 'work/wd/weightdecay0.001'
setproctitle.setproctitle(args.save)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
if os.path.exists(args.save):
shutil.rmtree(args.save)
os.makedirs(args.save, exist_ok=True)
trainLoader = torch.utils.data.DataLoader(dset.ImageFolder(
'/dataset/102flower/102flowers_20_10/train/',
transforms.Compose([
transforms.Scale(72),
transforms.RandomSizedCrop(64),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.434, 0.385, 0.296],
std=[0.285, 0.238, 0.262])
])),
batch_size=args.batchSz,
shuffle=True,
num_workers=2,
pin_memory=True)
valLoader = torch.utils.data.DataLoader(dset.ImageFolder(
'/dataset/102flower/102flowers_20_10/val/',
transforms.Compose([
transforms.Scale(72),
transforms.CenterCrop(64),
transforms.ToTensor(),
transforms.Normalize(mean=[0.434, 0.385, 0.296],
std=[0.285, 0.238, 0.262])
])),
batch_size=args.batchSz,
shuffle=False,
num_workers=2,
pin_memory=True)
net = resnet.ResNet18()
print(' + Number of params: {}'.format(
sum([p.data.nelement() for p in net.parameters()])))
if args.cuda:
net = net.cuda()
net = torch.nn.DataParallel(net,
device_ids=range(
torch.cuda.device_count()))
cudnn.benchmark = True
if args.opt == 'sgd':
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.wd)
elif args.opt == 'adam':
optimizer = optim.Adam(net.parameters(), weight_decay=1e-4)
elif args.opt == 'rmsprop':
optimizer = optim.RMSprop(net.parameters(), weight_decay=1e-4)
trainF = open(os.path.join(args.save, 'train.csv'), 'w')
testF = open(os.path.join(args.save, 'test.csv'), 'w')
for epoch in range(1, args.nEpochs + 1):
adjust_opt(args.opt, optimizer, args.lr, epoch)
train(args, epoch, net, trainLoader, optimizer, trainF)
test(args, epoch, net, valLoader, optimizer, testF)
torch.save(net, os.path.join(args.save, 'latest.pth'))
os.system('./plot.py {} &'.format(args.save))
trainF.close()
testF.close()
def return_accuracies(start_idxs, NUM_ROUND, NUM_QUERY, epoch, learning_rate,
datadir, data_name, feature):
torch.manual_seed(42)
torch.cuda.manual_seed(42)
np.random.seed(42)
random.seed(42)
torch.backends.cudnn.deterministic = True
if data_name in [
'dna', 'sklearn-digits', 'satimage', 'svmguide1', 'letter',
'shuttle', 'ijcnn1', 'sensorless', 'connect_4', 'sensit_seismic',
'usps', 'adult'
]:
fullset, valset, testset, num_cls = load_dataset_numpy_old(
datadir, data_name, feature=feature)
write_knndata_old(datadir, data_name, feature=feature)
elif data_name == 'cifar10':
fullset, valset, testset, num_cls = load_dataset_pytorch(
datadir, data_name)
# Validation Data set is 10% of the Entire Trainset.
validation_set_fraction = 0.1
num_fulltrn = len(fullset)
num_val = int(num_fulltrn * validation_set_fraction)
num_trn = num_fulltrn - num_val
trainset, validset = random_split(
fullset,
[num_trn, num_val]) #,generator=torch.Generator().manual_seed(42))
'''x_trn = fullset.data[trainset.indices]
y_trn = torch.from_numpy(np.array(fullset.targets)[trainset.indices].astype('float32'))
x_val = fullset.data[validset.indices]
y_val = torch.from_numpy(np.array(fullset.targets)[validset.indices])
x_tst = testset.data
y_tst = torch.from_numpy(np.array(testset.targets))'''
trn_batch_size = 128
val_batch_size = 1000
tst_batch_size = 1000
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=trn_batch_size,
shuffle=False,
pin_memory=True)
valloader = torch.utils.data.DataLoader(valset,
batch_size=val_batch_size,
shuffle=False,
sampler=SubsetRandomSampler(
validset.indices),
pin_memory=True)
testloader = torch.utils.data.DataLoader(testset,
batch_size=tst_batch_size,
shuffle=False,
pin_memory=True)
for batch_idx, (inputs, targets) in enumerate(trainloader):
if batch_idx == 0:
x_trn = inputs
y_trn = targets
else:
x_trn = torch.cat([x_trn, inputs], dim=0)
y_trn = torch.cat([y_trn, targets], dim=0)
for batch_idx, (inputs, targets) in enumerate(valloader):
if batch_idx == 0:
x_val = inputs
y_val = targets
else:
x_val = torch.cat([x_val, inputs], dim=0)
y_val = torch.cat([y_val, targets], dim=0)
for batch_idx, (inputs, targets) in enumerate(testloader):
if batch_idx == 0:
x_tst = inputs
y_tst = targets
else:
x_tst = torch.cat([x_tst, inputs], dim=0)
y_tst = torch.cat([y_tst, targets], dim=0)
#y_tst = y_tst.numpy().astype('float32')
#y_val = y_val.numpy().astype('float32')
y_trn = y_trn.numpy().astype('float32')
elif data_name in ['mnist', "fashion-mnist"]:
fullset, testset, num_cls = load_dataset_numpy_old(datadir,
data_name,
feature=feature)
write_knndata_old(datadir, data_name, feature=feature)
else:
fullset, valset, testset, num_cls = load_dataset_numpy(datadir,
data_name,
feature=feature)
write_knndata(datadir, data_name, feature=feature)
if data_name == 'mnist' or data_name == "fashion-mnist":
x_trn, y_trn = fullset.data, fullset.targets
x_tst, y_tst = testset.data, testset.targets
x_trn = x_trn.view(x_trn.shape[0], -1).numpy()
x_tst = x_tst.view(x_tst.shape[0], -1).numpy()
y_trn = y_trn.numpy()
y_tst = y_tst.numpy() #.float()
# Get validation data: Its 10% of the entire (full) training data
x_trn, x_val, y_trn, y_val = train_test_split(x_trn,
y_trn,
test_size=0.1,
random_state=42)
else:
if data_name != 'cifar10':
x_trn, y_trn = fullset
x_val, y_val = valset
x_tst, y_tst = testset
if data_name == 'cifar10':
handler = DataHandler3
else:
handler = CustomDataset_WithId
if data_name == 'cifar10':
args = {
'n_epoch': epoch,
'transform': None,
'loader_tr_args': {
'batch_size': 128
},
'loader_te_args': {
'batch_size': 1000
},
'optimizer_args': {
'lr': learning_rate
},
'transformTest': None
}
#transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
args['lr'] = learning_rate
else:
args = {
'transform': None,
'n_epoch': epoch,
'loader_tr_args': {
'batch_size': NUM_QUERY
},
'loader_te_args': {
'batch_size': 1000
},
'optimizer_args': {
'lr': learning_rate
},
'transformTest': None
}
args['lr'] = learning_rate
n_pool = len(y_trn)
n_val = len(y_val)
n_test = len(y_tst)
if data_name == 'cifar10':
net = resnet.ResNet18()
else:
net = mlpMod(x_trn.shape[1], num_cls,
100) #linMod(x_trn.shape[1], num_cls)
idxs_lb = np.zeros(n_pool, dtype=bool)
idxs_lb[start_idxs] = True
strategy = BadgeSampling(x_trn, y_trn, idxs_lb, net, handler, args)
strategy.train()
unlabled_acc = np.zeros(NUM_ROUND + 1)
tst_acc = np.zeros(NUM_ROUND + 1)
val_acc = np.zeros(NUM_ROUND + 1)
P = strategy.predict(x_tst, y_tst)
tst_acc[0] = 100.0 * P.eq(torch.tensor(y_tst)).sum().item() / n_test
print('\ttesting accuracy {}'.format(tst_acc[0]), flush=True)
#tst_acc[0] = 100.0 * P.eq(torch.tensor(y_tst)).sum().item()/ n_test
#print('\ttesting accuracy {}'.format(tst_acc[0]), flush=True)
P = strategy.predict(x_val, y_val)
val_acc[0] = 100.0 * P.eq(torch.tensor(y_val)).sum().item() / n_val
#idxs_unlabeled = (idxs_lb == False).nonzero().flatten().tolist()
u_x_trn = x_trn[~idxs_lb]
u_y_trn = y_trn[~idxs_lb]
P = strategy.predict(u_x_trn, u_y_trn)
unlabled_acc[0] = 100.0 * P.eq(
torch.tensor(u_y_trn)).sum().item() / len(u_y_trn)
for rd in range(1, NUM_ROUND + 1):
print('Round {}'.format(rd), flush=True)
# query
output = strategy.query(NUM_QUERY)
q_idxs = output
idxs_lb[q_idxs] = True
# report weighted accuracy
#corr = (strategy.predict(X_tr[q_idxs], torch.Tensor(Y_tr.numpy()[q_idxs]).long())).numpy() == Y_tr.numpy()[q_idxs]
# update
strategy.update(idxs_lb)
strategy.train()
# round accuracy
P = strategy.predict(x_tst, y_tst)
tst_acc[rd] = 100.0 * P.eq(torch.tensor(y_tst)).sum().item() / n_test
print(rd, '\ttesting accuracy {}'.format(tst_acc[rd]), flush=True)
P = strategy.predict(x_val, y_val)
val_acc[rd] = 100.0 * P.eq(torch.tensor(y_val)).sum().item() / n_val
#idxs_unlabeled = (idxs_lb == False).nonzero().flatten().tolist()
u_x_trn = x_trn[~idxs_lb]
u_y_trn = y_trn[~idxs_lb]
P = strategy.predict(u_x_trn, u_y_trn)
unlabled_acc[rd] = 100.0 * P.eq(
torch.tensor(u_y_trn)).sum().item() / len(u_y_trn)
#print(str(sum(idxs_lb)) + '\t' + 'unlabled data', len(u_y_trn),flush=True)
if sum(~strategy.idxs_lb) < NUM_QUERY:
sys.exit('too few remaining points to query')
return val_acc, tst_acc, unlabled_acc, np.arange(n_pool)[idxs_lb]
def run():
generator = Generator().to(device)
teacher = torch.load(opt.teacher_dir + 'teacher').to(device)
teacher.eval()
criterion = torch.nn.CrossEntropyLoss().to(device)
teacher = nn.DataParallel(teacher)
generator = nn.DataParallel(generator)
def kdloss(y, teacher_scores):
p = F.log_softmax(y, dim=1)
q = F.softmax(teacher_scores, dim=1)
l_kl = F.kl_div(p, q, size_average=False) / y.shape[0]
return l_kl
if opt.dataset == 'MNIST':
# Configure data loader
net = LeNet5Half().to(device)
net = nn.DataParallel(net)
data_test = MNIST(opt.data,
train=False,
transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
]))
data_test_loader = DataLoader(data_test, batch_size=64, num_workers=1, shuffle=False)
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr_G)
optimizer_S = torch.optim.Adam(net.parameters(), lr=opt.lr_S)
if opt.dataset != 'MNIST':
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
if opt.dataset == 'cifar10':
net = resnet.ResNet18().to(device)
net = nn.DataParallel(net)
data_test = CIFAR10(opt.data,
train=False,
transform=transform_test)
if opt.dataset == 'cifar100':
net = resnet.ResNet18(num_classes=100).to(device)
net = nn.DataParallel(net)
data_test = CIFAR100(opt.data,
train=False,
transform=transform_test)
data_test_loader = DataLoader(data_test, batch_size=opt.batch_size, num_workers=0)
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr_G)
optimizer_S = torch.optim.SGD(net.parameters(), lr=opt.lr_S, momentum=0.9, weight_decay=5e-4)
def adjust_learning_rate(optimizer, epoch, learing_rate):
if epoch < 800:
lr = learing_rate
elif epoch < 1600:
lr = 0.1*learing_rate
else:
lr = 0.01*learing_rate
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# ----------
# Training
# ----------
batches_done = 0
accr_best = 0
for epoch in range(opt.n_epochs):
total_correct = 0
avg_loss = 0.0
if opt.dataset != 'MNIST':
adjust_learning_rate(optimizer_S, epoch, opt.lr_S)
for i in range(120):
net.train()
z = Variable(torch.randn(opt.batch_size, opt.latent_dim)).to(device)
optimizer_G.zero_grad()
optimizer_S.zero_grad()
gen_imgs = generator(z)
outputs_T, features_T = teacher(gen_imgs, out_feature=True)
pred = outputs_T.data.max(1)[1]
loss_activation = -features_T.abs().mean()
loss_one_hot = criterion(outputs_T,pred)
softmax_o_T = torch.nn.functional.softmax(outputs_T, dim = 1).mean(dim = 0)
loss_information_entropy = (softmax_o_T * torch.log(softmax_o_T)).sum()
loss = loss_one_hot * opt.oh + loss_information_entropy * opt.ie + loss_activation * opt.a
loss_kd = kdloss(net(gen_imgs.detach()), outputs_T.detach())
loss += loss_kd
loss.backward()
optimizer_G.step()
optimizer_S.step()
if i == 1:
print ("[Epoch %d/%d] [loss_oh: %f] [loss_ie: %f] [loss_a: %f] [loss_kd: %f]" % (epoch, opt.n_epochs,loss_one_hot.item(), loss_information_entropy.item(), loss_activation.item(), loss_kd.item()))
with torch.no_grad():
for i, (images, labels) in enumerate(data_test_loader):
images = images.to(device)
labels = labels.to(device)
net.eval()
output = net(images)
avg_loss += criterion(output, labels).sum()
pred = output.data.max(1)[1]
total_correct += pred.eq(labels.data.view_as(pred)).sum()
avg_loss /= len(data_test)
print('Test Avg. Loss: %f, Accuracy: %f' % (avg_loss.data.item(), float(total_correct) / len(data_test)))
accr = round(float(total_correct) / len(data_test), 4)
if accr > accr_best:
torch.save(net,opt.output_dir + 'student')
torch.save(generator.state_dict(), opt.output_dir + "generator.pt")
accr_best = accr
def get_trained_model(self):
"""
Trains a model on all labelled samples until validation loss stops decreasing
:return: trained model
"""
# Initialize model and optimization tools.
if self.model_type == "simplecnn":
model = SimpleCNN()
elif self.model_type == "resnet18":
model = resnet.ResNet18()
else:
print(f"Unrecognized model type {self.model_type}!")
raise SystemExit
model.to(self.device)
model_path = Path(f"model_{uuid.uuid4().hex}")
while model_path.exists():
model_path = Path(f"model_{uuid.uuid4().hex}")
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.RMSprop(model.parameters(),
lr=LEARNING_RATE,
momentum=MOMENTUM,
weight_decay=5e-4)
min_validation_loss = np.inf
best_epoch = -1
print(f"Training on {self.labelled_idx.shape[0]} samples.")
for epoch in range(MAX_EPOCHS):
# training
model.train()
running_loss = 0.0
for i, batch in enumerate(self.train_loader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = batch[0].to(self.device), batch[1].to(
self.device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 20 == 19: # print every 20 mini-batches
print(
f'[{epoch}: {i + 1}/{len(self.train_loader)}] loss: {running_loss / 20:.3f}'
)
running_loss = 0.0
# validation
model.eval()
with torch.no_grad():
average_validation_loss = 0.0
for i, batch in enumerate(self.validation_loader, 0):
inputs, labels = batch[0].to(self.device), batch[1].to(
self.device)
outputs = model(inputs)
average_validation_loss += criterion(outputs, labels)
average_validation_loss /= (VALIDATION_SET_SIZE /
EVALUATION_BATCH_SIZE)
print(f"Validation loss: {average_validation_loss:.3f}")
# check for validation loss improvement
if average_validation_loss < min_validation_loss:
min_validation_loss = average_validation_loss
best_epoch = epoch
# create checkpoint of model with lowest validation loss so far
torch.save(model.state_dict(), model_path)
print(
f"Training completed. Best epoch was {best_epoch} with validation loss {min_validation_loss:.3f}"
)
model.load_state_dict(torch.load(model_path))
model_path.unlink()
return model, criterion
import torch
import resnet
import torch.backends.cudnn as cudnn
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = resnet.ResNet18().to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
checkpoint = torch.load('ResNet18_ckpt.pth')
net.load_state_dict(checkpoint['net'])
test = torch.rand([1, 3, 32, 32])
net.eval()
out = net(test)
print(out)
testset = torchvision.datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform_test)
testloader = torch.utils.data.DataLoader(testset,
batch_size=100,
shuffle=False,
num_workers=2)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
# Model
print('==> Building model..')
net = RN.ResNet18()
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
if args.ce == True:
criterion = nn.CrossEntropyLoss()
save_path = './checkpoint/CrossEntropy.bin'
print("Use CrossEntropy")
else:
criterion = LabelSmoothingCrossEntropy()
save_path = './checkpoint/LabelSmoothing.bin'
print("Use Label Smooting")
transforms.Resize((64, 64)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = datasets.CIFAR10(root=args.data_path,
train=True,
download=True,
transform=transform)
train_loader = torch.utils.data.DataLoader(trainset,
batch_size=args.batch_size,
num_workers=args.workers)
generator = model.Generator(100, [1024, 512, 256, 128], 3).train()
fid_model = resnet.ResNet18().eval()
mod_state = torch.load(args.fid_model_path,
map_location=lambda storage, loc: storage)
fid_model.load_state_dict(mod_state['model_state'])
if args.cuda:
generator = generator.cuda()
if not os.path.isfile('./test_data_statistics.p'):
testset = datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform)
test_loader = torch.utils.data.DataLoader(testset,
batch_size=1000,
shuffle=False,
num_workers=2)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
train_acc_matrix = np.zeros(X.shape)
train_loss_matrix = np.zeros(X.shape)
test_acc_matrix = np.zeros(X.shape)
test_loss_matrix = np.zeros(X.shape)
for i in range(chunk_id * int(X.shape[0] / 4),
(chunk_id + 1) * int(X.shape[0] / 4)):
print('Now on i ' + str(i))
for j in range(X.shape[1]):
net_on_grid = resnet.ResNet18()
net_on_grid.load_state_dict(
torch.load('results/curve_find/15418861362468774/net_i_' + str(i) +
'_j_' + str(j) + '.t7'))
net_on_grid = net_on_grid.to(device)
net_on_grid.train()
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
outputs = net_on_grid(inputs)
net_on_grid.eval()
for loaderid, loader in enumerate([testloader, trainloader]):
if loaderid == 0:
batch_size=64,
num_workers=1,
shuffle=False)
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr_G)
optimizer_S = torch.optim.Adam(net.parameters(), lr=opt.lr_S)
if opt.dataset != 'MNIST':
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
if opt.dataset == 'cifar10':
net = resnet.ResNet18().cuda()
net = nn.DataParallel(net)
data_test = CIFAR10(opt.data,
train=False,
transform=transform_test,
download=True)
if opt.dataset == 'cifar100':
net = resnet.ResNet18(num_classes=100).cuda()
net = nn.DataParallel(net)
data_test = CIFAR100(opt.data, train=False, transform=transform_test)
data_test_loader = DataLoader(data_test,
batch_size=opt.batch_size,
num_workers=0)
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr_G)
yearpolars[0] = 1
# Determine the polar regression values.
trainy = 2. * ((years - lb) / (ub - lb)) - 1
trainloader.dataset.years = trainy
trainloader.dataset.oldyears = years.copy()
# Other prototype loading.
stylepolars = np.load("dat/prototypes/sgd/prototypes-46d-46c.npy")
stylepolars = torch.from_numpy(stylepolars).float()
allpolars = [yearpolars, stylepolars]
# Network type.
if args.network == "std":
model = convnet.Std(args.output_dims, None)
elif args.network == "resnet16":
model = resnet.ResNet18(args.output_dims, None)
elif args.network == "resnet32":
model = resnet.ResNet34(args.output_dims, None)
model = model.to(device)
# To CUDA.
if args.multigpu == 1:
model = torch.nn.DataParallel(model.cuda())
else:
model = model.to(device)
# Network parameters.
optimname = args.optimizer
lr = args.learning_rate
momentum = args.momentum
decay = args.decay
