def _get_specified_model(model):
if model == 'ResNet18':
return ResNet18()
elif model == 'ResNet34':
return ResNet34()
elif model == 'PreActResNet18':
return PreActResNet18()
elif model == 'PreActResNet34':
return PreActResNet34()
elif model == 'WideResNet28':
return WideResNet28()
elif model == 'WideResNet34':
return WideResNet34()
else:
return ResNet18()
def init_model():
print('==> Building model..'+args.model)
if args.model == 'ResNet18':
net = ResNet18(num_classes=num_classes)
elif args.model == 'MLP':
# 4-layer MLP
input_dim = 3072 if ('CIFAR' in args.dataset) else 784
width = args.width
net = torch.nn.Sequential(OrderedDict([
('flatten', torch.nn.Flatten()),
('linear0', torch.nn.Linear(input_dim, width)),
('relu0', torch.nn.ReLU()),
('linear1', torch.nn.Linear(width, width)),
('relu1', torch.nn.ReLU()),
('linear2', torch.nn.Linear(width, width)),
('relu2', torch.nn.ReLU()),
('linear3', torch.nn.Linear(width, num_classes))]))
elif args.model == 'DenseNet':
net = densenet_cifar(num_classes=num_classes)
elif args.model == 'MobileNetV2':
net = MobileNetV2(num_classes=num_classes)
elif args.model == 'ResNet20_FIXUP':
net = fixup_resnet20(num_classes=num_classes)
else:
raise ValueError('shitty args.model name')
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
return net
def configure_model(self):
"""
:return:
"""
arch: str = self.hparams.arch
batch_norm = self.hparams.batch_norm
dataset: str = self.hparams.dataset
hidden_layers: int = self.hparams.hidden_layers
hidden_size: int = self.hparams.hidden_size
if arch == 'mlp':
if dataset == 'mnist':
return MLP(input_size=784,
hidden_size=hidden_size,
num_hidden_layers=hidden_layers,
batch_norm=batch_norm)
elif dataset == 'cifar10':
return MLP(hidden_size=hidden_size,
num_hidden_layers=hidden_layers,
batch_norm=batch_norm)
else:
raise ValueError('invalid dataset specification!')
elif arch == 'alexnet':
return AlexNet()
elif arch == 'vgg11':
return VGG(vgg_name='VGG11')
elif arch == 'vgg13':
return VGG(vgg_name='VGG13')
elif arch == 'resnet18':
return ResNet18()
elif arch == 'resnet34':
return ResNet34()
else:
raise ValueError('Unsupported model!')
def get_model(name: str):
name = name.lower()
if name == 'vgg11':
return VGG('VGG11')
elif name == 'resnet18':
return ResNet18()
elif name == 'resnet34':
return ResNet34()
elif name == 'resnet50':
return ResNet50()
def get_model(dataset_type):
if dataset_type == 'mnist':
model = MNISTNet().to(device)
elif dataset_type == 'cifar10':
# model = CIFAR10Net().to(device)
model = CNN9Layer(num_classes=10, input_shape=3).to(device)
# model = ResNet18(num_classes=10).to(device)
elif dataset_type == 'cifar100':
# model = CNN9Layer(num_classes=100, input_shape=3).to(device)
model = ResNet18(num_classes=100).to(device)
return model
def _preact_resnet18(msda='fmix', pretrained=False, repeat=0, *args, **kwargs):
from models import ResNet18
model = ResNet18(*args, **kwargs)
if pretrained:
state = load_state_dict_from_url(
f'http://marc.ecs.soton.ac.uk/pytorch-models/cifar10/preact-resnet18/cifar10_preact_resnet18_{msda}_{repeat}.pt',
progress=True)
model.load_state_dict(state)
return model
def main():
if not os.path.exists(SAVE_DIR):
os.makedirs(SAVE_DIR)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Using PyTorch device : {}'.format(device.upper()))
model = ResNet18().to(device)
train_loader, val_loader, test_loader = utils.get_cifar10_data_loaders(n_train=N_TRAIN, \
n_val=N_VAL, n_test=N_TEST)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-1, momentum=0.9)
xent = nn.CrossEntropyLoss()
step = 0
n_epochs = 50
f = open(SAVE_DIR + LOSS_NAME, 'w')
f.truncate(0)
f.write('train_step, train_loss\n')
f.close()
MODEL_SAVE_PATH = SAVE_DIR + MODEL_NAME
assert model.training is True
for i in range(n_epochs):
for j, (images, labels) in enumerate(train_loader):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = xent(outputs, labels)
loss.backward()
optimizer.step()
if step % 16 == 0:
with open(SAVE_DIR + LOSS_NAME, 'a') as f:
f.write('{}, {}\n'.format(step, loss.item()))
step += 1
utils.progress(
j + 1, len(train_loader),
'Batch [{}/{}] Epoch [{}/{}]'.format(j + 1, len(train_loader),
i + 1, n_epochs))
torch.save(model.state_dict(), MODEL_SAVE_PATH.format(i))
print('Training Complete')
model.eval()
correct = 0
for j, (images, labels) in enumerate(test_loader):
images, labels = images.to(device), labels.to(device)
logits = model(images)
correct += (torch.max(logits, 1)[-1] == labels).sum().item()
utils.progress(j + 1, len(test_loader),
'Batch [{}/{}]'.format(j + 1, len(test_loader)))
print('Test Accuracy = {}%'.format(float(correct) * 100.0 / N_TEST))
print('Done')
def load_model(config):
if config['model'] == 'ResNet18':
model = ResNet18(color_channel=config['color_channel'])
elif config['model'] == 'VGG11':
model = VGG('VGG11', color_channel=config['color_channel'])
elif config['model'] == 'VGG13':
model = VGG('VGG13', color_channel=config['color_channel'])
else:
print('wrong model option')
model = None
model_path = config['dir_path'] + '/models/' + config['data'] + '_' + config['model'] + '_t1=' + \
str(config['t1']) + '_R=' + config['R'] + "_" + config['fixed'] + '.pt'
model.load_state_dict(torch.load(model_path))
model.cuda()
return model
def get_model(argu, device):
if argu.arch == "c11":
model = c11()
elif argu.arch == "resnet18":
model = ResNet18()
elif argu.arch == "resnet8":
model = ResNet8()
if argu.pretrain:
model.load_state_dict(torch.load(argu.pretrain, map_location=device))
model.to(device)
print("\n *** pretrain model loaded: "+ argu.pretrain + " *** \n")
return model.to(device)
def get_model_for_attack(model_name):
if model_name == 'model1':
model = ResNet34()
load_w(model, "./models/weights/resnet34.pt")
elif model_name == 'model2':
model = ResNet18()
load_w(model, "./models/weights/resnet18_AT.pt")
elif model_name == 'model3':
model = SmallResNet()
load_w(model, "./models/weights/res_small.pth")
elif model_name == 'model4':
model = WideResNet34()
pref = next(model.parameters())
model.load_state_dict(
filter_state_dict(
torch.load("./models/weights/trades_wide_resnet.pt",
map_location=pref.device)))
elif model_name == 'model5':
model = WideResNet()
load_w(model, "./models/weights/wideres34-10-pgdHE.pt")
elif model_name == 'model6':
model = WideResNet28()
pref = next(model.parameters())
model.load_state_dict(
filter_state_dict(
torch.load('models/weights/RST-AWP_cifar10_linf_wrn28-10.pt',
map_location=pref.device)))
elif model_name == 'model_vgg16bn':
model = vgg16_bn(pretrained=True)
elif model_name == 'model_resnet18_imgnet':
model = resnet18(pretrained=True)
elif model_name == 'model_inception':
model = inception_v3(pretrained=True)
elif model_name == 'model_vitb':
from mnist_vit import ViT, MegaSizer
model = MegaSizer(
ImageNetRenormalize(ViT('B_16_imagenet1k', pretrained=True)))
elif model_name.startswith('model_hub:'):
_, a, b = model_name.split(":")
model = torch.hub.load(a, b, pretrained=True)
model = Cifar10Renormalize(model)
elif model_name.startswith('model_mnist:'):
_, a = model_name.split(":")
model = torch.load('mnist.pt')[a]
elif model_name.startswith('model_ex:'):
_, a = model_name.split(":")
model = torch.load(a)
return model
def val_our_model():
model = ResNet18()
model.load_state_dict(torch.load('./checkpoints/' + 'resnet' + '_' + '0607_14_30_57.pth'))
model.eval()
model.cuda()
label_extra = Label_Extraction('label1.csv')
labels = label_extra.get_data()
print('Running model on the validation data:')
val_data = Cancer_Pic('data1', labels, train=False)
val(model, val_data)
labels = label_extra.get_data()
print('Running model on the test data:')
test_data = Cancer_Pic('data1', labels, test=True)
val(model, test_data)
def build_model(config):
if config['model'] == 'ResNet18':
model = ResNet18(color_channel=config['color_channel'])
elif config['model'] == 'VGG11':
model = VGG('VGG11', color_channel=config['color_channel'])
elif config['model'] == 'VGG13':
model = VGG('VGG13', color_channel=config['color_channel'])
else:
print('wrong model option')
model = None
loss_function = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=config['lr'],
momentum=config['momentum'],
weight_decay=config['weight_decay'])
return model, loss_function, optimizer
def build_model(device, model_name, num_classes=10):
"""构建模型：vgg、vggnonorm、resnet、preactresnet、googlenet、densenet、
resnext、mobilenet、mobilenetv2、dpn、shufflenetg2、senet、shufflenetv2
:param device: 'cuda' if you have a GPU, 'cpu' otherwise
:param model_name: One of the models available in the folder 'models'
:param num_classes: 10 or 100 depending on the chosen dataset
:return: The model architecture
"""
print('==> Building model..')
model_name = model_name.lower()
if model_name == 'vgg':
net = VGG('VGG19', num_classes=num_classes)
elif model_name == 'vggnonorm':
net = VGG('VGG19', num_classes=num_classes, batch_norm=False)
elif model_name == 'resnet':
net = ResNet18(num_classes=num_classes)
elif model_name == 'preactresnet':
net = PreActResNet18()
elif model_name == 'googlenet':
net = GoogLeNet()
elif model_name == 'densenet':
net = DenseNet121()
elif model_name == 'resnext':
net = ResNeXt29_2x64d()
elif model_name == 'mobilenet':
net = MobileNet()
elif model_name == 'mobilenetv2':
net = MobileNetV2()
elif model_name == 'dpn':
net = DPN92()
elif model_name == 'shufflenetg2':
net = ShuffleNetG2()
elif model_name == 'senet':
net = SENet18()
elif model_name == 'shufflenetv2':
net = ShuffleNetV2(1)
else:
raise ValueError('Error: the specified model is incorrect ({})'.format(model_name))
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
return net
def get_model_for_defense(model_name):
if 'WRN28' in model_name:
model = WideResNet28()
elif 'PreActRN18' in model_name:
model = PreActResNet18()
elif 'PreActRN34' in model_name:
model = PreActResNet34()
elif 'RN18' in model_name:
model = ResNet18()
elif 'RN34' in model_name:
model = ResNet34()
else:
raise ValueError(
f'Unsupported model name: {model_name}. Check your spelling!')
checkpoint = torch.load(f"./models/weights/{model_name}.pt")
model.load_state_dict(checkpoint['model'])
return model
def get_model(argu, device):
if argu.arch == "c11":
model = c11()
elif argu.arch == "resnet18":
model = ResNet18()
elif argu.arch == "resnet8":
model = ResNet8()
elif argu.arch == "c12":
model = c12()
else:
raise NotImplementedError("model not included")
if argu.pretrain:
model.load_state_dict(torch.load(argu.pretrain, map_location=device))
model.to(device)
print("\n *** pretrain model loaded: "+ argu.pretrain + " *** \n")
return model.to(device)
def get_model_for_attack(model_name):
if model_name=='model1':
model = ResNet34()
model.load_state_dict(torch.load("models/weights/resnet34.pt"))
elif model_name=='model2':
model = ResNet18()
model.load_state_dict(torch.load('models/weights/resnet18_AT.pt'))
elif model_name=='model3':
model = SmallResNet()
model.load_state_dict(torch.load('models/weights/res_small.pth'))
elif model_name=='model4':
model = WideResNet34()
model.load_state_dict(filter_state_dict(torch.load('models/weights/trades_wide_resnet.pt')))
elif model_name=='model5':
model = WideResNet()
model.load_state_dict(torch.load('models/weights/wideres34-10-pgdHE.pt'))
elif model_name=='model6':
model = WideResNet28()
model.load_state_dict(filter_state_dict(torch.load('models/weights/RST-AWP_cifar10_linf_wrn28-10.pt')))
return model
def test_model_epsilon_range(model_name,
attack,
epsilon_range,
num_iters,
test_effort=True,
ideal_SNR=0.6,
device="cpu"):
if args.test_mnist:
model_to_test = MNIST_Net()
else:
model_to_test = ResNet18()
if gpu_enabled:
model_to_test.load_state_dict(torch.load(input_directory + model_name))
else:
model_to_test.load_state_dict(
torch.load(input_directory + model_name,
map_location={'cuda:0': 'cpu'}))
model_to_test.to(device)
lst_efforts = []
for epsilon in epsilon_range:
lst_efforts.append(
test_model(model_to_test,
attack,
epsilon,
num_iters,
ideal_SNR,
test_effort=test_effort,
device=device))
print 'finished testing model: {} with epsilon: {}'.format(
model_name, str(epsilon))
print lst_efforts
np.savetxt(output_directory + model_name + extension,
np.array(lst_efforts),
fmt='%f')
return lst_efforts
def main():
'''model = models.resnet18(pretrained = False)
model.fc = nn.Linear(512, 4000)'''
model = ResNet18(4000)
if (param['resume'] == True):
if (param['load checkpoint from N_way_classification'] == True):
checkPointPath = '../checkpoints/N_way_classification' + '/epoch' + str(
param['resume_from'])
checkpoint = torch.load(checkPointPath)
model.load_state_dict(checkpoint['model_state_dict'])
'''modules = list(model.children())[:-1]
model = nn.Sequential(*modules)'''
else:
checkPointPath = param['checkPointPath'] + '/epoch' + str(
param['resume_from'])
checkpoint = torch.load(checkPointPath)
model.load_state_dict(checkpoint['model_state_dict'])
model.to(DEVICE)
optimizer = Adam(model.parameters(), lr=0.0001)
scheduler = lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.95)
criterion = TripletLoss(margin=1, img_per_person=param['img_per_person'])
criterion.to(DEVICE)
data_transform = transforms.Compose(
[transforms.Resize((224, 224)),
transforms.ToTensor()])
train_root = "../classification_data/train_data"
dev_path = '../verification_pairs_val.txt'
test_path = '../verification_pairs_test.txt'
train_data = train_dataset(train_root, param['img_per_person'],
data_transform)
dev_data = dev_dataset(dev_path, data_transform)
test_data = dev_dataset(test_path, data_transform)
train_loader = DataLoader(train_data,
batch_size=param['train_batch_size'],
pin_memory=True,
shuffle=True)
dev_loader = DataLoader(dev_data,
batch_size=param['dev_batch_size'],
pin_memory=True,
shuffle=False)
test_loader = DataLoader(test_data,
batch_size=param['dev_batch_size'],
pin_memory=True,
shuffle=False)
print("start training")
start_epoch = param['resume_from'] + 1
torch.cuda.empty_cache()
for epoch in range(start_epoch, start_epoch + param['nepochs']):
train(model, train_loader, criterion, optimizer, epoch)
path = param['checkPointPath'] + "/epoch" + str(epoch)
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict()
}, path)
scheduler.step()
auc = validation(model, dev_loader, param['thresh'])
print("validation auc is: ", auc)
test(model, test_loader)
def train_gan(opt):
os.makedirs(os.path.join(
opt.savingroot, opt.dataset,
str(opt.p1 * 100) + '%complementary/' + str(opt.p1) + '_chkpts'),
exist_ok=True)
if not os.path.exists(
os.path.join(
opt.savingroot, opt.data_r, 'data',
'processed/training' + str(opt.p1) + str(opt.p2) + '.pt')):
generate_c_data(opt)
else:
print('data exits')
#Build network
if opt.data_r == 'MNIST':
netd_c = LeNet5(opt.ndf, opt.nc, num_classes=10).cuda()
netd_glj = D_MNIST(opt.ndf, opt.nc, num_classes=10).cuda()
elif opt.data_r == 'CIFAR10':
netd_c = ResNet18().cuda(
) #DPN92().cuda()#D_CIFAR10(opt.ndf, opt.nc, num_classes=10).cuda()
optd_c = optim.SGD(
netd_c.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4
) #optim.Adam(netd_c.parameters(), lr=0.0002, betas=(0.5, 0.999),weight_decay=5e-4) #
optd_glj = optim.SGD(
netd_glj.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4
) #optim.Adam(netd_c.parameters(), lr=0.0002, betas=(0.5, 0.999),weight_decay=5e-4) #
print('training_start')
step1 = 0
step2 = 0
acc1 = []
acc2 = []
# configure(os.path.join(opt.savingroot,opt.dataset,str(opt.p1 * 100) + '%complementary/' + str(opt.p1) + '/logs'), flush_secs=5)
test_loader = torch.utils.data.DataLoader(CIFAR10_Complementary(
os.path.join(opt.savingroot, opt.data_r, 'data'),
train=False,
unlabel=False,
transform=transforms.Compose([
transforms.Resize(opt.image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])),
batch_size=128,
num_workers=2)
fixed = embed_z(opt)
dataset = CIFAR10_Complementary(os.path.join(opt.savingroot, opt.data_r,
'data'),
train=True,
unlabel=False,
transform=tsfm,
p1=opt.p1,
p2=opt.p2)
loader = torch.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=2,
worker_init_fn=np.random.seed)
loader1 = loader
loader2 = loader
#####################unlabel data ##################################################
unlabeldataset = CIFAR10_Complementary(
os.path.join(opt.savingroot, opt.data_r, 'data'),
train=False,
unlabel=True,
transform=transforms.Compose([
transforms.Resize(opt.image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]),
p1=opt.p1,
p2=opt.p2)
unlabel_loader = torch.utils.data.DataLoader(unlabeldataset,
batch_size=128,
num_workers=2)
unlabel_loader1 = unlabel_loader
unlabel_loader2 = unlabel_loader
##################### ##################3###############################
choose_dataset10, choose_dataset11, choose_dataset20, choose_dataset21 = torch.Tensor(
[]).cuda(), torch.Tensor([]).cuda(), torch.Tensor(
[]).cuda(), torch.Tensor([]).cuda()
co_data1, co_data2 = torch.Tensor([]).cuda(), torch.Tensor([]).cuda()
for epoch in range(opt.num_epoches):
print('Epoch:{:03d}'.format(epoch))
if epoch % int(opt.num_epoches / 3) == 0 and epoch != 0:
for param_group in optd_c.param_groups:
param_group['lr'] = param_group['lr'] / 10
print(param_group['lr'])
for param_group in optd_glj.param_groups:
param_group['lr'] = param_group['lr'] / 10
print(param_group['lr'])
#################model1
step1 = train_c(epoch,
netd_c,
optd_c,
loader1,
step1,
opt,
co_data=co_data1) #
accur1 = test_acc(netd_c, test_loader, opt, False) #
'''
if epoch>5:
choose_data2,new_unlabel_data1=test_acc(netd_c, unlabel_loader1 , opt,True)
unlabeldataset=CIFAR10_Complementary(os.path.join(opt.savingroot,opt.data_r,'data'), train=False,unlabel=True,transform=transforms.Compose([
transforms.Resize(opt.image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]),p1=opt.p1,p2=opt.p2,new_unlabeldata=new_unlabel_data1)
new_unlabel_loader1 = torch.utils.data.DataLoader( unlabeldataset,batch_size=128, num_workers=2)
unlabel_loader1=new_unlabel_loader1###############################
print(choose_data2[0].size(),choose_data2[1].size(),new_unlabel_data1[0].size(),new_unlabel_data1[1].size())
'''
acc1.append(accur1)
#################model2
step2 = train_c(epoch,
netd_glj,
optd_glj,
loader2,
step2,
opt,
co_data=co_data2) #
accur2 = test_acc(netd_glj, test_loader, opt, False)
'''
if epoch>5:
choose_data1,new_unlabel_data2=test_acc(netd_glj, unlabel_loader2 , opt,True)
unlabeldataset=CIFAR10_Complementary(os.path.join(opt.savingroot,opt.data_r,'data'), train=False,unlabel=True,transform=transforms.Compose([
transforms.Resize(opt.image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]),p1=opt.p1,p2=opt.p2,new_unlabeldata=new_unlabel_data2)
new_unlabel_loader2 = torch.utils.data.DataLoader( unlabeldataset,batch_size=128, num_workers=2)
unlabel_loader2=new_unlabel_loader2
print(choose_data1[0].size(),choose_data1[1].size(),new_unlabel_data2[0].size(),new_unlabel_data2[1].size())
'''
acc2.append(accur2)
'''
if epoch>5:
####################co_data1
choose_dataset10 = torch.cat([choose_dataset10,choose_data1[0]],dim=0)
choose_dataset11 = torch.cat([choose_dataset11,choose_data1[1]],dim=0)
print(choose_dataset10.size(),choose_dataset11.size())
co_data1=[choose_dataset10,choose_dataset11]
#new_traindata1=CIFAR10_Complementary(os.path.join(opt.savingroot, opt.data_r, 'data'),unlabel=False, transform=tsfm, p1=opt.p1,p2=opt.p2,co_data=co_data1)
####################co_data2
choose_dataset20 = torch.cat([choose_dataset20,choose_data2[0]],dim=0)
choose_dataset21 = torch.cat([choose_dataset21,choose_data2[1]],dim=0)
print(choose_dataset20.size(),choose_dataset21.size())
co_data2=[choose_dataset20,choose_dataset21]
#new_traindata2=CIFAR10_Complementary(os.path.join(opt.savingroot, opt.data_r, 'data'),unlabel=False, transform=tsfm, p1=opt.p1,p2=opt.p2,co_data=co_data2)
####################
#new_train_loader1= torch.utils.data.DataLoader(new_traindata2, batch_size=opt.batch_size, shuffle=True, num_workers=2,worker_init_fn=np.random.seed)
#loader1=new_train_loader1
#new_train_loader2= torch.utils.data.DataLoader(new_traindata1, batch_size=opt.batch_size, shuffle=True, num_workers=2,worker_init_fn=np.random.seed)
#loader2=new_train_loader2
'''
f = open(
os.path.join(opt.savingroot, opt.dataset,
str(opt.p1 * 100) + '%complementary/' +
'LeNet5_acc1.txt'), 'w')
for cont in acc1:
f.writelines(str(cont) + '\n')
f.close()
f = open(
os.path.join(opt.savingroot, opt.dataset,
str(opt.p1 * 100) + '%complementary/' +
'DMNIST_acc2.txt'), 'w')
for cont in acc2:
f.writelines(str(cont) + '\n')
f.close()
if args.wmtrain:
print('Loading watermark images')
wmloader = getwmloader(args.wm_path, args.wm_batch_size, args.wm_lbl)
# create the model
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.exists(args.load_path), 'Error: no checkpoint found!'
checkpoint = torch.load(args.load_path)
net = checkpoint['net']
acc = checkpoint['acc']
start_epoch = checkpoint['epoch']
else:
print('==> Building model..')
net = ResNet18(num_classes=n_classes)
net = net.to(device)
# support cuda
if device == 'cuda':
print('Using CUDA')
print('Parallel training on {0} GPUs.'.format(torch.cuda.device_count()))
net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
# loading wm examples
if args.wmtrain:
print("WM acc:")
# model select
if modelName == 'AlexNet':
from models.AlexNet import *
model = AlexNet()
elif modelName == "SqueezeNet":
from models.SqueezeNet import *
model = SqueezeNet()
elif modelName == "VGG16":
from models.VGG16 import *
model = VGG16()
elif modelName == "GoogLeNet":
from models.GoogLeNet import *
model = GoogLeNet()
elif modelName == "ResNet18":
from models.ResNet18 import *
model = ResNet18()
elif modelName == "DenseNet121":
from models.DenseNet import *
model = DenseNet121()
elif modelName == "MobileNet":
from models.MobileNet import *
model = MobileNet()
else:
model = None
if model is not None:
model = model.to(Device)
else:
raise Exception("model not found")
def create_torch_model():
model = ResNet18()
model = load_checkpoint(model, pretrained='checkpoint/ckpt.pth')
return model
def main():
'''model = models.resnet18(pretrained = False)
model.fc = nn.Linear(512, 4000)'''
model = ResNet18(4000)
model.to(DEVICE)
optimizer = SGD(model.parameters(),
lr=0.15,
momentum=0.9,
weight_decay=5e-5)
if (param['resume'] == True):
print("loading from checkpoint {}".format(param['resume_from']))
checkPointPath = param['checkPointPath'] + '/epoch' + str(
param['resume_from'])
checkpoint = torch.load(checkPointPath)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.to(DEVICE)
print("finish loading")
scheduler = lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.95)
criterion = nn.CrossEntropyLoss()
criterion.to(DEVICE)
batch_size = 10 if DEVICE == 'cuda' else 1
data_transform = transforms.Compose(
[transforms.Resize((224, 224)),
transforms.ToTensor()])
train_dataset = datasets.ImageFolder(
root='../classification_data/train_data/', transform=data_transform)
val_dataset = datasets.ImageFolder(root='../classification_data/val_data',
transform=data_transform)
verfication_dev_dataset = MyVerificationDataset(
'../verification_pairs_val.txt', data_transform)
verfication_test_dataset = MyVerificationDataset(
'../verification_pairs_test.txt', data_transform)
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True)
dev_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
verification_dev_loader = DataLoader(verfication_dev_dataset,
batch_size=batch_size,
shuffle=False)
verification_test_loader = DataLoader(verfication_test_dataset,
batch_size=batch_size,
shuffle=False)
start_epoch = param['resume_from'] + 1
torch.cuda.empty_cache()
acc = validation(model, dev_loader)
auc = verification_dev(model, verification_dev_loader)
print("start training")
for epoch in range(start_epoch, start_epoch + param['nepochs']):
train(model, train_loader, criterion, optimizer, epoch)
path = param['checkPointPath'] + "/epoch" + str(epoch)
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict()
}, path)
scheduler.step()
acc = validation(model, dev_loader)
auc = verification_dev(model, verification_dev_loader)
print("auc is: ", auc)
if __name__ == "__main__":
args = parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
gpu_num = max(len(args.gpu_id.split(',')), 1)
model_name = 'resnet18'
log_dir = "logs/%s_%s" % (time.strftime("%b%d-%H%M",
time.localtime()), model_name)
check_mkdir(log_dir)
log = Logger(log_dir + '/train.log')
log.print(args)
device = torch.device('cuda')
model = ResNet18().to(device)
model = nn.DataParallel(model, device_ids=[i for i in range(gpu_num)])
train_loader, test_loader = prepare_cifar(args.batch_size,
args.test_batch_size)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
best_epoch, best_robust_acc = 0, 0.
for e in range(args.epoch):
adjust_learning_rate(optimizer, e)
train_acc, train_robust_acc, loss = train_adv_epoch(
model, args, train_loader, device, optimizer, e)
if e % 3 == 0 or (e >= 74 and e <= 80):
# get data loader
train_loader, val_loader, test_loader, clean_sample_idx, noisy_sample_idx, dataset_len = create_dataloader(
args.dataset, args.dataset_path, args.noise_type, args.noise_rate)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if args.dataset == 'mnist':
model = MNISTNet().to(device)
elif args.dataset == 'cifar10':
# model = CIFAR10Net().to(device)
model = CNN9Layer(num_classes=10, input_shape=3).to(device)
# model = ResNet18(num_classes=10).to(device)
elif args.dataset == 'cifar100':
# model = CNN9Layer(num_classes=100, input_shape=3).to(device)
model = ResNet18(num_classes=100).to(device)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.l2_reg)
train_loss_lst, val_loss_lst = [], []
train_acc_lst, val_acc_lst = [], []
clean_mean_loss_lst, noisy_mean_loss_lst = [], []
clean_min_loss_lst, noisy_min_loss_lst = [], []
clean_max_loss_lst, noisy_max_loss_lst = [], []
sample_loss = np.zeros(dataset_len)
elif args.dataset == 'gtsrb':
transform_train = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])
model = GTSRBNet(dim=1).to(device)
elif args.dataset == 'cifar10':
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
])
model = ResNet18(dim=3).to(device)
optimizer = optim.Adam(model.parameters(), lr=0.0001)
Xy_train = MyDataset(dataset=args.dataset,
dtype=args.dtype,
train=True,
transform=transform_train)
Xy_test = MyDataset(dataset=args.dataset,
dtype=args.dtype,
train=False,
transform=transform_test)
train_loader = torch.utils.data.DataLoader(Xy_train,
shuffle=True,
batch_size=args.batch_size,
train_transform = get_transforms(224)
train_dataset = ImageNetDataset(data_path, 'train', train_transform)
train_dataloader = DataLoader(train_dataset,
batch_size=256,
shuffle=True,
num_workers=6)
if torch.cuda.is_available():
device = torch.device("cuda:0")
print("Running on the GPU")
else:
device = torch.device("cpu")
print("Running on the CPU")
model = ResNet18(1000).to(device)
loss_fun = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=0.1,
momentum=0.9,
weight_decay=1e-4)
nepochs = 100
for epoch in range(nepochs):
epoch_start = time.time()
print('')
print('Epoch ' + str(epoch))
for i, data in enumerate(train_dataloader, 0):
inputs, labels = data
inputs = inputs.to(device)
def return_model(mode, **kwargs):
if mode == 'logistic':
solver = kwargs.get('solver', 'liblinear')
n_jobs = kwargs.get('n_jobs', None)
max_iter = kwargs.get('max_iter', 5000)
model = LogisticRegression(solver=solver,
n_jobs=n_jobs,
max_iter=max_iter,
random_state=666,
multi_class='auto')
elif mode == 'Tree':
model = DecisionTreeClassifier(random_state=666)
elif mode == 'RandomForest':
n_estimators = kwargs.get('n_estimators', 50)
model = RandomForestClassifier(n_estimators=n_estimators,
random_state=666)
elif mode == 'GB':
n_estimators = kwargs.get('n_estimators', 50)
model = GradientBoostingClassifier(n_estimators=n_estimators,
random_state=666)
elif mode == 'AdaBoost':
n_estimators = kwargs.get('n_estimators', 50)
model = AdaBoostClassifier(n_estimators=n_estimators, random_state=666)
elif mode == 'SVC':
kernel = kwargs.get('kernel', 'rbf')
model = SVC(kernel=kernel, random_state=666)
elif mode == 'LinearSVC':
model = LinearSVC(loss='hinge', random_state=666)
elif mode == 'GP':
model = GaussianProcessClassifier(random_state=666)
elif mode == 'KNN':
n_neighbors = kwargs.get('n_neighbors', 5)
model = KNeighborsClassifier(n_neighbors=n_neighbors)
elif mode == 'NB':
model = MultinomialNB()
elif mode == 'linear':
model = LinearRegression(random_state=666)
elif mode == 'ridge':
alpha = kwargs.get('alpha', 1.0)
model = Ridge(alpha=alpha, random_state=666)
elif mode == 'ResNet':
model = ResNet18(num_classes=kwargs.get('num_classes', 10))
elif 'conv' in mode:
tf.reset_default_graph()
address = kwargs.get('address', 'weights/conv')
hidden_units = kwargs.get('hidden_layer_sizes', [20])
activation = kwargs.get('activation', 'relu')
weight_decay = kwargs.get('weight_decay', 1e-4)
learning_rate = kwargs.get('learning_rate', 0.001)
max_iter = kwargs.get('max_iter', 1000)
early_stopping = kwargs.get('early_stopping', 10)
warm_start = kwargs.get('warm_start', False)
batch_size = kwargs.get('batch_size', 256)
kernel_sizes = kwargs.get('kernel_sizes', [5])
strides = kwargs.get('strides', [5])
channels = kwargs.get('channels', [1])
validation_fraction = kwargs.get('validation_fraction', 0.)
global_averaging = kwargs.get('global_averaging', 0.)
optimizer = kwargs.get('optimizer', 'sgd')
if mode == 'conv':
model = CShapNN(mode='classification',
batch_size=batch_size,
max_epochs=max_iter,
learning_rate=learning_rate,
weight_decay=weight_decay,
validation_fraction=validation_fraction,
early_stopping=early_stopping,
optimizer=optimizer,
warm_start=warm_start,
address=address,
hidden_units=hidden_units,
strides=strides,
global_averaging=global_averaging,
kernel_sizes=kernel_sizes,
channels=channels,
random_seed=666)
elif mode == 'conv_reg':
model = CShapNN(mode='regression',
batch_size=batch_size,
max_epochs=max_iter,
learning_rate=learning_rate,
weight_decay=weight_decay,
validation_fraction=validation_fraction,
early_stopping=early_stopping,
optimizer=optimizer,
warm_start=warm_start,
address=address,
hidden_units=hidden_units,
strides=strides,
global_averaging=global_averaging,
kernel_sizes=kernel_sizes,
channels=channels,
random_seed=666)
elif 'NN' in mode:
solver = kwargs.get('solver', 'sgd')
hidden_layer_sizes = kwargs.get('hidden_layer_sizes', (20, ))
if isinstance(hidden_layer_sizes, list):
hidden_layer_sizes = list(hidden_layer_sizes)
activation = kwargs.get('activation', 'relu')
learning_rate_init = kwargs.get('learning_rate', 0.001)
max_iter = kwargs.get('max_iter', 5000)
early_stopping = kwargs.get('early_stopping', False)
warm_start = kwargs.get('warm_start', False)
if mode == 'NN':
model = MLPClassifier(solver=solver,
hidden_layer_sizes=hidden_layer_sizes,
activation=activation,
learning_rate_init=learning_rate_init,
warm_start=warm_start,
max_iter=max_iter,
early_stopping=early_stopping)
if mode == 'NN_reg':
model = MLPRegressor(solver=solver,
hidden_layer_sizes=hidden_layer_sizes,
activation=activation,
learning_rate_init=learning_rate_init,
warm_start=warm_start,
max_iter=max_iter,
early_stopping=early_stopping)
else:
raise ValueError("Invalid mode!")
return model
import torch
import torchvision
from models import ResNet18
# loading the model
path = "pretrained/pretrained88.pth"
m = ResNet18()
m.load_state_dict(torch.load(path))
# load test set
transform = torchvision.transforms.ToTensor()
testset = torchvision.datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform)
testloader = torch.utils.data.DataLoader(testset,
batch_size=128,
shuffle=False,
num_workers=0)
# evaluate clean accuracy (=90.84)
# since images are unnormalized (between 0 and 1), we want the model to
# perform the normalization internally
m.set_normalize(True)
m.eval()
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = m(images)
batch_size=100,
shuffle=False,
num_workers=0)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
logger = Logger(os.path.join('./checkpoint', args.exp_name, 'logs'))
with open(os.path.join('./checkpoint', args.exp_name, 'args.pkl'), "wb") as f:
pickle.dump(args, f)
# Model
print('==> Building model..')
# net = VGG('VGG19')
net = ResNet18()
checkpoint = torch.load('res/mobile_net_py3.pth')
net.load_state_dict(checkpoint, strict=False)
net.cuda()
net.eval()
path_dims = [64, 64, 128, 128, 256, 256, 512, 512]
switch_vec = args.switch
print('Switch vector {}'.format(switch_vec))
sup_net = SupervisorNetwork(path_dims)
models_dir = os.path.join('./checkpoint', args.exp_name, 'models')
load_epoch = 0