def main():
''' Main function, flow of program '''
# Model
model = resnet18()
# Running architecture (GPU or CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Image loader
test_loader = get_dataloader(params['test_file'], params['img_size'],\
params['batch_size'], params['data_mean'], params['data_std'])
# Creates the criterion (loss function)
criterion = nn.CrossEntropyLoss()
# Weights Load Up
weights_file = glob.glob(params['weights_path'] + '/*.pth')[0]
checkpoint = torch.load(weights_file)
model.load_state_dict(checkpoint['model_state_dict'])
print('Model Loaded!\nAccuracy: {:.4}\nLoss: {:.4}\nSensitivity: {:.4}\nSpecificity: {:.4}'\
.format(checkpoint['accuracy'], checkpoint['loss'],\
checkpoint['sensitivity'], checkpoint['specificity']))
# Create folder for weights
pathlib.Path(params['report_path']).mkdir(parents=True, exist_ok=True)
# Run test and creates a report
test_report(model, test_loader, criterion, device, params['report_path'])
def run(args: DictConfig) -> None:
assert torch.cuda.is_available()
torch.manual_seed(args.seed)
n_classes = args.get(args.dataset).n_classes
classifier = resnet18(n_classes=n_classes).to(args.device)
logger.info('Base classifier resnet18: # parameters {}'.format(
cal_parameters(classifier)))
data_dir = hydra.utils.to_absolute_path(args.data_dir)
train_data = get_dataset(data_name=args.dataset,
data_dir=data_dir,
train=True,
crop_flip=True)
test_data = get_dataset(data_name=args.dataset,
data_dir=data_dir,
train=False,
crop_flip=False)
train_loader = DataLoader(dataset=train_data,
batch_size=args.n_batch_train,
shuffle=True)
test_loader = DataLoader(dataset=test_data,
batch_size=args.n_batch_test,
shuffle=False)
if args.inference:
save_name = 'resnet18_wd{}.pth'.format(args.weight_decay)
classifier.load_state_dict(
torch.load(save_name, map_location=lambda storage, loc: storage))
loss, acc = run_epoch(classifier, test_loader, args)
logger.info('Inference, test loss: {:.4f}, Acc: {:.4f}'.format(
loss, acc))
else:
train(classifier, train_loader, test_loader, args)
def run(args: DictConfig) -> None:
assert torch.cuda.is_available()
torch.manual_seed(args.seed)
n_classes = args.get(args.dataset).n_classes
rep_size = args.get(args.dataset).rep_size
margin = args.get(args.dataset).margin
classifier = resnet18(n_classes=n_classes).to(args.device)
sdim = SDIM(disc_classifier=classifier,
n_classes=n_classes,
rep_size=rep_size,
mi_units=args.mi_units,
margin=margin).to(args.device)
base_dir = hydra.utils.to_absolute_path('logs/sdim/{}'.format(
args.dataset))
save_name = 'SDIM_resnet18{}.pth'.format(suffix_dict[args.base_type])
sdim.load_state_dict(
torch.load(os.path.join(base_dir, save_name),
map_location=lambda storage, loc: storage))
if args.sample_likelihood:
sample_cases(sdim, args)
else:
pgd_attack(sdim, args)
def get_model():
# Get model from config
if config.model == "resnet18":
model = models.resnet18(pretrained=config.pretrained)
elif config.model == "resnet34":
model = models.resnet34(pretrained=config.pretrained)
elif config.model == 'resnet50':
model = models.resnet50(pretrained=config.pretrained)
elif config.model == "resnet101":
model = models.resnet101(pretrained=config.pretrained)
elif config.model == "resnet152":
model = models.resnet152(pretrained=config.pretrained)
elif config.model == "resnext50_32x4d":
model = models.resnet34(pretrained=config.pretrained)
elif config.model == 'resnext101_32x8d':
model = models.resnet50(pretrained=config.pretrained)
elif config.model == "wide_resnet50_2":
model = models.resnet101(pretrained=config.pretrained)
elif config.model == "wide_resnet101_2":
model = models.resnet152(pretrained=config.pretrained)
else:
raise ValueError('%s not supported'.format(config.model))
# Initialize fc layer
(in_features, out_features) = model.fc.in_features, model.fc.out_features
model.fc = torch.nn.Linear(in_features, out_features)
return model
def _set_models(self):
f_net = resnet18(**self.options.f_config)
g_nets = [bagnet18(**self.options.g_config)
for _ in range(self.options.n_g_nets)]
self.model = ReBiasModels(f_net, g_nets)
self.evaluator = ImageNetEvaluator(device=self.device)
def prep_model(init_model=None):
model = models.resnet18(pretrained=False)
num_frts = model.fc.in_features
model.fc = torch.nn.Linear(num_frts, 2)
if init_model is None:
return model
model.load_state_dict(torch.load(init_model))
return model
def load_pretrained_model(args):
""" load pretrained base discriminative classifier."""
n_classes = args.get(args.dataset).n_classes
classifier = resnet18(n_classes=n_classes).to(args.device)
if not args.inference:
save_name = 'resnet18{}.pth'.format(suffix_dict[args.base_type])
base_dir = 'logs/base/{}'.format(args.dataset)
path = hydra.utils.to_absolute_path(base_dir)
classifier.load_state_dict(torch.load(os.path.join(path, save_name)))
return classifier
def show_save_tensor():
import torch
from torchvision import utils
import torchvision.models as models
from matplotlib import pyplot as plt
def vis_tensor(tensor, ch=0, all_kernels=False, nrow=8, padding=2):
'''
ch: channel for visualization
allkernels: all kernels for visualization
'''
n, c, h, w = tensor.shape
if all_kernels:
tensor = tensor.view(n * c, -1, w, h)
elif c != 3:
tensor = tensor[:, ch, :, :].unsqueeze(dim=1)
rows = np.min((tensor.shape[0] // nrow + 1, 64))
grid = utils.make_grid(tensor,
nrow=nrow,
normalize=True,
padding=padding)
# plt.figure(figsize=(nrow,rows))
plt.imshow(grid.numpy().transpose((1, 2, 0))) #CHW HWC
def save_tensor(tensor,
filename,
ch=0,
all_kernels=False,
nrow=8,
padding=2):
n, c, h, w = tensor.shape
if all_kernels:
tensor = tensor.view(n * c, -1, w, h)
elif c != 3:
tensor = tensor[:, ch, :, :].unsqueeze(dim=1)
utils.save_image(tensor,
filename,
nrow=nrow,
normalize=True,
padding=padding)
vgg = models.resnet18(pretrained=True)
mm = vgg.double()
filters = mm.modules
body_model = [i for i in mm.children()][0]
# layer1 = body_model[0]
layer1 = body_model
tensor = layer1.weight.data.clone()
vis_tensor(tensor)
save_tensor(tensor, 'test.png')
plt.axis('off')
plt.ioff()
plt.show()
def model_generator(device, opt):
model = resnet18()
if opt.load_model_path:
model.load_state_dict(t.load(opt.load_model_path))
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Linear(model.fc.in_features, opt.nums_of_classes)
if t.cuda.device_count() > 1:
model = nn.DataParallel(model, device_ids=opt.device_ids)
model.to(device)
return model
def test():
model = models.resnet18()
print model.layer1[0].conv1.weight.data
print model.layer1[
0].conv1.__class__ #<class 'torch.nn.modules.conv.Conv2d'>
print model.layer1[0].conv1.kernel_size
input = torch.autograd.Variable(torch.randn(20, 16, 50, 100))
print input.size()
print np.prod(input.size())
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--input-file',
'-i',
type=str,
default=None,
help='Path to an input file.')
parser.add_argument(
'--model',
'-m',
type=str,
choices=['resnet18Fer2013', 'resnet18KDEF', 'resnet18Dartmouth'],
default='resnet18Fer2013',
help='Model name.')
args = parser.parse_args()
input_file_name = args.input_file
model_name = args.model
img = Image.open(input_file_name)
if model_name == 'resnet18Fer2013':
net = resnet18()
weights = 'resnet18Fer2013.pth'
transform = get_transform(48)
elif model_name == 'resnet18KDEF':
net = resnet18()
weights = 'resnet18KDEF.pth'
transform = get_transform(224)
elif model_name == 'resnet18Dartmouth':
net = resnet18()
weights = 'resnet18Dartmouth.pth'
transform = get_transform(224)
path = Path.joinpath(Path(), 'weights', weights)
net.load_state_dict(torch.load(path, map_location=torch.device('cpu')))
net.eval()
img_tensor = transform(img).unsqueeze(0)
with torch.no_grad():
softmax = torch.nn.Softmax(dim=1)
preds = softmax(net(img_tensor)).numpy().ravel()
for i in range(len(CLASSES)):
print('{:>8}: {:5.2f}%'.format(CLASSES[i], 100 * preds[i]))
def get_imagenet_models(model_name):
if model_name == 'model_vgg16bn':
from models import vgg16_bn
model = vgg16_bn(pretrained=True)
elif model_name == 'model_resnet18_imgnet':
from models import resnet18
model = resnet18(pretrained=True)
elif model_name == 'model_inception':
from models import inception_v3
model = inception_v3(pretrained=True)
else:
raise ValueError(f'Buggya no model named {model_name}')
# print(f'Model: {model_name}')
return model
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 test():
import models
DEVICE = torch.device('cuda:{}'.format(0))
net = models.resnet18(pretrained=False)
net.fc = nn.Linear(512, 257)
net.to(DEVICE)
criterion = nn.CrossEntropyLoss().to(DEVICE)
x = torch.ones(5, 3, 224, 224).to(DEVICE)
labels = torch.ones(5).type(torch.long).to(DEVICE)
attack = SpatialTransfromedAttack()
adv_inp = attack(x, labels, net, criterion)
print(adv_inp.size())
def get_model(model_name, input_shape, nclasses, args):
n = input_shape[-1]
if model_name == 'vtcnn2':
return models.vt_cnn2(input_shape, nclasses)
elif model_name == 'resnet18-outer':
model, preprocess = models.resnet18((n, n, 3), nclasses, args.dropout)
def f(x):
return preprocess(256 * preprocessing.preprocess_outer(x))
return model, f
elif model_name == 'resnet18-gasf':
model, preprocess = models.resnet18((n, n, 3), nclasses, args.dropout)
def f(x):
return preprocess(256 * preprocessing.preprocess_gasf(x))
return model, f
elif model_name == 'resnet18-gadf':
model, preprocess = models.resnet18((n, n, 3), nclasses, args.dropout)
def f(x):
return preprocess(256 * preprocessing.preprocess_gadf(x))
return model, f
elif model_name == 'resnet18-mtf':
model, preprocess = models.resnet18((n, n, 3), nclasses, args.dropout)
def f(x):
return preprocess(256 * preprocessing.preprocess_mtf(x))
return model, f
elif model_name == 'resnet18-rplot':
model, preprocess = models.resnet18((n, n, 3), nclasses, args.dropout)
def f(x):
return preprocess(256 * preprocessing.preprocess_rplot(x))
return model, f
elif model_name == 'resnet18-noisy-outer':
model, preprocess = models.resnet18((n, n, 3), nclasses, args.dropout)
def f(x):
return preprocess(256 * preprocessing.preprocess_noisy_outer(x))
return model, f
else:
raise ValueError("Unknown model %s", model_name)
def print_layers_num():
# resnet = models.resnet18()
resnet = models.resnet18()
def foo(net):
childrens = list(net.children())
if not childrens:
if isinstance(net, torch.nn.Conv2d):
print ' '
#可以用来统计不同层的个数
# net.register_backward_hook(print)
return 1
count = 0
for c in childrens:
count += foo(c)
return count
print(foo(resnet))
def main(args):
model = resnet18(num_classes=2)
# model = models.densenet121(num_classes = 2)
model = torch.nn.DataParallel(model, device_ids=list(range(
args.num_gpus))).cuda()
cudnn.benchmark = True
if args.resume:
if os.path.isfile(args.resume):
print("=> Loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
state_dict = checkpoint['state_dict']
model.load_state_dict(state_dict)
print("=> Loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
raise Exception("=> No checkpoint found at '{}'".format(
args.resume))
tf = TestDataset(val_dir, args)
num_images = len(val_dir)
# dataset = datasets.ImageFolder(root=args.root_path, transform = data_transforms )
# tf = TrainDataset(train_dir, args)
train_loader = DataLoader(tf,
batch_size=args.batch_size,
shuffle=args.shuffle,
num_workers=args.num_workers,
pin_memory=True)
print("Start testing ...")
train(train_loader, model, args)
# elapsed_time = time.time() - start_time
# print('epoch time ' +str(time.strftime("%H:%M:%S", time.gmtime(elapsed_time))) + ', remaining time ' +str(time.strftime("%H:%M:%S", time.gmtime(elapsed_time*(args.num_epochs - epoch)))))
# save models
# if (epoch >= args.particular_epoch and epoch % args.save_epochs_steps == 0) or epoch % args.save_epochs_steps_before == 0 :
# if epoch % args.save_epochs_steps == 0:
# save_checkpoint({'epoch': epoch, 'state_dict': model.state_dict(), 'opt_dict': optimizer.state_dict()}, epoch, args)
print("Testing Done")
def main():
model = resnet18(hparams)
model.to(device)
optimizer = optim.SGD(model.parameters(), lr=1e-4)
multi_gpu_model = nn.DataParallel(model, device_ids=gpu_devices)
objects = {
"model": multi_gpu_model,
"optim": optimizer,
"device": gpu_id,
"loss_fn": loss_fn
}
scheduler = Scheduler(config, objects)
# call the training loop
for epoch in range(config["total_epochs"]):
scheduler.train()
scheduler.test()
def test():
test_config = {
'model_path': 'checkpoints/resnet18_final_checkpoint.pth',
'visualize': False
}
print(test_config)
device = t.device("cuda" if t.cuda.is_available() else "cpu")
opt = DefaultConfig()
model = resnet18().eval()
model.fc = nn.Linear(model.fc.in_features, opt.nums_of_classes)
model.load_state_dict(t.load(os.path.join(os.getcwd(), test_config['model_path'])))
model.to(device)
print('Model loaded!')
# data
test_data = BaldDataset(opt.test_data_root)
test_dataloader = DataLoader(dataset=test_data,
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.num_workers,
collate_fn=test_data.customized_collate_fn,
drop_last=False)
print('Data loaded!')
correct = 0
print('Start inference.....')
with t.no_grad():
f = open(opt.result_file, 'w+', encoding='utf-8', newline='')
csv_writer = csv.writer(f)
csv_writer.writerow(['image', 'detection_result', 'ground_truth'])
for inputs, target, img_names in tqdm(test_dataloader):
inputs = inputs.to(device)
target = target.to(device)
output = model(inputs)
pred = t.relu(output).data.max(1)[1]
correct += pred.eq(target.data).sum()
for img_name, dr, gt in zip(img_names, pred, target.data):
csv_writer.writerow([img_name, dr.item(), gt.item()])
test_acc = correct.cpu().detach().numpy() * 1.0 / len(test_dataloader.dataset)
f.close()
print("test_acc: %.3f \n" % test_acc)
print('Classification results were saved in path {}.'.format(opt.result_file))
def build_model(model_name, num_classes, pretrain):
if model_name == 'resnet50':
net = resnet50(num_classes=num_classes, pretrain=pretrain)
elif model_name == 'resnet18':
net = resnet18(num_classes=num_classes, pretrain=pretrain)
elif model_name == 'resnet34':
net = resnet34(num_classes=num_classes, pretrain=pretrain)
elif model_name == 'resnet101':
net = resnet101(num_classes=num_classes, pretrain=pretrain)
elif model_name == 'resnet152':
net = resnet152(num_classes=num_classes, pretrain=pretrain)
elif model_name == 'resnet50se':
net = resnet50se(num_classes=num_classes, pretrain=pretrain)
elif model_name == 'resnet50dilated':
net = resnet50_dilated(num_classes=num_classes, pretrain=pretrain)
elif model_name == 'resnet50dcse':
net = resnet50_dcse(num_classes=num_classes, pretrain=pretrain)
else:
print('wait a minute')
return net
def test(**kwargs):
# opt.parse(kwargs)
# import ipdb;
# ipdb.set_trace()
# configure model
#model = getattr(models, opt.model)().eval()
model=models.resnet18(pretrained=True)
if opt.load_model_path:
model.load(opt.load_model_path)
if opt.use_gpu: model.cuda()
# data
test_data = DogCat(opt.test_data_root,test=True)
test_dataloader = DataLoader(test_data,batch_size=opt.batch_size,shuffle=True,num_workers=opt.num_workers)
#print(test_dataloader)
print("come")
results = []
for ii,(data,path) in tqdm(enumerate(test_dataloader)):
input = t.autograd.Variable(data,volatile = True)
#print("data is {}".format(data))
print("path is {}".format(path))
if opt.use_gpu: input = input.cuda()
score = model(input)
print("score is :{}".format(score))
probability = t.nn.functional.softmax(score)[:,0].data.tolist()
print("probability is :".format(probability))
label = score.max(dim = 1)[1].data.tolist()
print("Label is {}".format(label))
batch_results = [(path_,probability_) for path_,probability_ in zip(path,probability) ]
results += batch_results
write_csv(results,opt.result_file)
print(results)
val_data = DogCat(opt.test_data_root,train=False)
val_dataloader = DataLoader(val_data,opt.batch_size,
shuffle=False,num_workers=opt.num_workers)
val_cm, val_accuracy = val(model, val_dataloader)
print("end val function")
print(val_cm.value(),val_accuracy)
print("test function end")
return results
def load_model(model_name):
if model_name == 'CNN_raw':
model = models.ConvNet()
model.weight_init(0, 0.02)
elif model_name == 'CNN_CBAM':
model = models.ConvNet_CBAM()
model.weight_init(0, 0.02)
elif model_name == 'ResNet18_raw':
model = models.resnet18()
elif model_name == 'ResNet18_CBAM':
model = models.resnet18_CBAM()
elif model_name == 'ResNet34_raw':
model = models.resnet34()
elif model_name == 'ResNet34_CBAM':
model = models.resnet34_CBAM()
elif model_name == 'ResNet50_raw':
model = models.resnet50()
elif model_name == 'ResNet50_CBAM':
model = models.resnet50_CBAM(num_classes=200)
else:
raise RuntimeError('Unknown model type!')
return model
def main():
''' Main function, flow of program '''
# To stablish a seed for all the project
seed_everything(parms['seed'])
# Model
model = resnet18()
# Running architecture (GPU or CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Images Loaders
train_loader = get_aug_dataloader(params['train_file'], params['img_size'],\
params['batch_size'], params['data_mean'], params['data_std'])
val_loader = get_dataloader(params['val_file'], params['img_size'],\
params['batch_size'], params['data_mean'], params['data_std'])
# Creates the criterion (loss function)
criterion = nn.CrossEntropyLoss()
# Creates optimizer (Changes the weights based on loss)
if params['optimizer'] == 'ADAM':
optimizer = torch.optim.Adam(model.parameters(),
lr=params['lear_rate'])
elif params['optimizer'] == 'SGD':
optimizer = torch.optim.SGD(model.parameters(),
lr=params['lear_rate'],
momentum=0.9)
# Create folder for weights
pathlib.Path(params['weights_path']).mkdir(parents=True, exist_ok=True)
# Training and Validation for the model
train_validate(model, train_loader, val_loader, optimizer,\
criterion, device, params['epochs'], params['save_criteria'],\
params['weights_path'])
def get_model_for_attack(model_name):
if model_name == 'model_vgg16bn':
model = vgg16_bn(pretrained=True)
elif model_name == 'model_resnet18':
model = resnet18(pretrained=True)
elif model_name == 'model_inceptionv3':
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)
else:
raise ValueError(f'Model f{model_name} does not exist.')
return model
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--resume',
type=str,
default='../exps/tradeoff.eps8/checkpoint.pth.tar')
parser.add_argument('-d', type=int, default=3)
args = parser.parse_args()
net_name = args.resume.split('/')[-2]
print(net_name)
path = os.path.join('../SmoothRes', net_name)
if not os.path.exists(path):
os.mkdir(path)
net = models.resnet18(pretrained=False)
net.fc.out_features = 200
net.load_state_dict(torch.load(args.resume)['state_dict'])
DEVICE = torch.device('cuda:{}'.format(args.d))
net.to(DEVICE)
dl = create_test_dataset(32)
#xz_test(dl, 1,net, DEVICE)
#test_model(net, dl)
#l1_for_without_smooth(net, dl, DEVICE)
#l1_for_with_smooth(net, dl, DEVICE)
get_result(net, dl, DEVICE, net_name)
def print_autograd_graph():
from graphviz import Digraph
import torch
from torch.autograd import Variable
def make_dot(var, params=None):
""" Produces Graphviz representation of PyTorch autograd graph
Blue nodes are the Variables that require grad, orange are Tensors
saved for backward in torch.autograd.Function
Args:
var: output Variable
params: dict of (name, Variable) to add names to node that
require grad (TODO: make optional)
"""
if params is not None:
#assert all(isinstance(p, Variable) for p in params.values())
param_map = {id(v): k for k, v in params.items()}
node_attr = dict(style='filled',
shape='box',
align='left',
fontsize='12',
ranksep='0.1',
height='0.2')
dot = Digraph(node_attr=node_attr, graph_attr=dict(size="12,12"))
seen = set()
def size_to_str(size):
return '(' + (', ').join(['%d' % v for v in size]) + ')'
def add_nodes(var):
if var not in seen:
if torch.is_tensor(var):
dot.node(str(id(var)),
size_to_str(var.size()),
fillcolor='orange')
elif hasattr(var, 'variable'):
u = var.variable
#name = param_map[id(u)] if params is not None else ''
#node_name = '%s\n %s' % (name, size_to_str(u.size()))
node_name = '%s\n %s' % (param_map.get(id(
u.data)), size_to_str(u.size()))
dot.node(str(id(var)), node_name, fillcolor='lightblue')
else:
dot.node(str(id(var)), str(type(var).__name__))
seen.add(var)
if hasattr(var, 'next_functions'):
for u in var.next_functions:
if u[0] is not None:
dot.edge(str(id(u[0])), str(id(var)))
add_nodes(u[0])
if hasattr(var, 'saved_tensors'):
for t in var.saved_tensors:
dot.edge(str(id(t)), str(id(var)))
add_nodes(t)
add_nodes(var.grad_fn)
return dot
from torchvision import models
torch.manual_seed(1)
inputs = torch.randn(1, 3, 224, 224)
model = models.resnet18(pretrained=False)
y = model(Variable(inputs))
#print(y)
g = make_dot(y, params=model.state_dict())
g.view()
def main(args):
if args.checkpoint == '':
args.checkpoint = "checkpoints/ctw1500_%s_bs_%d_ep_%d" % (
args.arch, args.batch_size, args.n_epoch)
if args.pretrain:
if 'synth' in args.pretrain:
args.checkpoint += "_pretrain_synth"
else:
args.checkpoint += "_pretrain_ic17"
print('checkpoint path: %s' % args.checkpoint)
print('init lr: %.8f' % args.lr)
print('schedule: ', args.schedule)
sys.stdout.flush()
if not os.path.isdir(args.checkpoint):
os.makedirs(args.checkpoint)
kernel_num = 7
min_scale = 0.4
start_epoch = 0
data_loader = CTW1500Loader(is_transform=True,
img_size=args.img_size,
kernel_num=kernel_num,
min_scale=min_scale)
#train_loader = ctw_train_loader(data_loader, batch_size=args.batch_size)
if args.arch == "resnet50":
model = models.resnet50(pretrained=True, num_classes=kernel_num)
elif args.arch == "resnet101":
model = models.resnet101(pretrained=True, num_classes=kernel_num)
elif args.arch == "resnet152":
model = models.resnet152(pretrained=True, num_classes=kernel_num)
#resnet18 and 34 didn't inplement pretrained
elif args.arch == "resnet18":
model = models.resnet18(pretrained=False, num_classes=kernel_num)
elif args.arch == "resnet34":
model = models.resnet34(pretrained=False, num_classes=kernel_num)
elif args.arch == "mobilenetv2":
model = models.resnet152(pretrained=True, num_classes=kernel_num)
elif args.arch == "mobilenetv3large":
model = models.mobilenetv3_large(pretrained=False,
num_classes=kernel_num)
elif args.arch == "mobilenetv3small":
model = models.mobilenetv3_small(pretrained=False,
num_classes=kernel_num)
optimizer = tf.keras.optimizers.SGD(learning_rate=args.lr,
momentum=0.99,
decay=5e-4)
title = 'CTW1500'
if args.pretrain:
print('Using pretrained model.')
assert os.path.isfile(
args.pretrain), 'Error: no checkpoint directory found!'
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Learning Rate', 'Train Loss', 'Train Acc.', 'Train IOU.'])
elif args.resume:
print('Resuming from checkpoint.')
model.load_weights(args.resume)
logger = Logger(os.path.join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print('Training from scratch.')
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Learning Rate', 'Train Loss', 'Train Acc.', 'Train IOU.'])
for epoch in range(start_epoch, args.n_epoch):
optimizer = get_new_optimizer(args, optimizer, epoch)
print(
'\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.n_epoch, optimizer.get_config()['learning_rate']))
train_loader = ctw_train_loader(data_loader,
batch_size=args.batch_size)
train_loss, train_te_acc, train_ke_acc, train_te_iou, train_ke_iou = train(train_loader, model, dice_loss,\
optimizer, epoch)
model.save_weights('%s%s' % (args.checkpoint, '/model_tf/weights'))
logger.append([
optimizer.get_config()['learning_rate'], train_loss, train_te_acc,
train_te_iou
])
logger.close()
def train_net(args):
torch.manual_seed(7)
np.random.seed(7)
checkpoint = args.checkpoint
start_epoch = 0
best_acc = float('-inf')
writer = SummaryWriter()
epochs_since_improvement = 0
# Initialize / load checkpoint
if checkpoint is None:
if args.network == 'r18':
model = resnet18(args)
elif args.network == 'r34':
model = resnet34(args)
elif args.network == 'r50':
model = resnet50(args)
elif args.network == 'r101':
model = resnet101(args)
elif args.network == 'r152':
model = resnet152(args)
elif args.network == 'mobile':
from mobilenet_v2 import MobileNetV2
model = MobileNetV2()
else:
raise TypeError('network {} is not supported.'.format(
args.network))
metric_fc = ArcMarginModel(args)
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=args.lr,
momentum=args.mom,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.Adam([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=args.lr,
weight_decay=args.weight_decay)
else:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
epochs_since_improvement = checkpoint['epochs_since_improvement']
model = checkpoint['model']
metric_fc = checkpoint['metric_fc']
optimizer = checkpoint['optimizer']
model = nn.DataParallel(model)
metric_fc = nn.DataParallel(metric_fc)
# Move to GPU, if available
model = model.to(device)
metric_fc = metric_fc.to(device)
# Loss function
if args.focal_loss:
criterion = FocalLoss(gamma=args.gamma).to(device)
else:
criterion = nn.CrossEntropyLoss().to(device)
# Custom dataloaders
train_dataset = ArcFaceDataset('train')
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=num_workers)
# Epochs
for epoch in range(start_epoch, args.end_epoch):
# Decay learning rate if there is no improvement for 2 consecutive epochs, and terminate training after 10
if epochs_since_improvement == 10:
break
if epochs_since_improvement > 0 and epochs_since_improvement % 2 == 0:
checkpoint = 'BEST_checkpoint.tar'
checkpoint = torch.load(checkpoint)
model = checkpoint['model']
metric_fc = checkpoint['metric_fc']
optimizer = checkpoint['optimizer']
adjust_learning_rate(optimizer, 0.5)
# One epoch's training
train_loss, train_top1_accs = train(train_loader=train_loader,
model=model,
metric_fc=metric_fc,
criterion=criterion,
optimizer=optimizer,
epoch=epoch)
lr = optimizer.param_groups[0]['lr']
print('\nCurrent effective learning rate: {}\n'.format(lr))
# print('Step num: {}\n'.format(optimizer.step_num))
writer.add_scalar('model/train_loss', train_loss, epoch)
writer.add_scalar('model/train_accuracy', train_top1_accs, epoch)
writer.add_scalar('model/learning_rate', lr, epoch)
if epoch % 5 == 0:
# One epoch's validation
megaface_acc = megaface_test(model)
writer.add_scalar('model/megaface_accuracy', megaface_acc, epoch)
# Check if there was an improvement
is_best = megaface_acc > best_acc
best_acc = max(megaface_acc, best_acc)
if not is_best:
epochs_since_improvement += 1
print("\nEpochs since last improvement: %d\n" %
(epochs_since_improvement, ))
else:
epochs_since_improvement = 0
# Save checkpoint
save_checkpoint(epoch, epochs_since_improvement, model, metric_fc,
optimizer, best_acc, is_best)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', default=0, type=int)
parser.add_argument('--model_arch',
default='conv4',
choices=['conv4', 'resnet10', 'resnet18'],
type=str)
# parser.add_argument('--attention', action='store_true')
parser.add_argument('--start_epoch', default=1, type=int)
parser.add_argument('--num_epoch', default=90, type=int)
parser.add_argument('--learning_rate', default=0.01, type=float)
parser.add_argument('--scheduler_milestones', nargs='+', type=int)
parser.add_argument('--alpha', default=0.5, type=float)
parser.add_argument('--model_saving_rate', default=30, type=int)
parser.add_argument('--train', action='store_true')
parser.add_argument('--support_groups', default=10000, type=int)
parser.add_argument('--evaluate', action='store_true')
parser.add_argument('--evaluation_rate', default=10, type=int)
parser.add_argument('--model_dir', default=None, type=str)
parser.add_argument('--checkpoint', action='store_true')
parser.add_argument('--normalize', action='store_true')
parser.add_argument('--save_settings', action='store_true')
parser.add_argument('--layer', default=4, type=int)
args = parser.parse_args()
device = torch.device(f'cuda:{args.gpu}')
model_arch = args.model_arch
# attention = args.attention
learning_rate = args.learning_rate
alpha = args.alpha
start_epoch = args.start_epoch
num_epoch = args.num_epoch
model_saving_rate = args.model_saving_rate
toTrain = args.train
toEvaluate = args.evaluate
evaluation_rate = args.evaluation_rate
checkpoint = args.checkpoint
normalize = args.normalize
scheduler_milestones = args.scheduler_milestones
save_settings = args.save_settings
support_groups = args.support_groups
# ------------------------------- #
# Generate folder
# ------------------------------- #
if checkpoint:
model_dir = f'./training_models/{args.model_dir}'
else:
model_dir = f'./training_models/{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}'
os.makedirs(model_dir)
# ------------------------------- #
# Config logger
# ------------------------------- #
train_logger = setup_logger('train_logger', f'{model_dir}/train.log')
result_logger = setup_logger('result_logger', f'{model_dir}/result.log')
# overview_logger = setup_logger('overview_logger', f'./overview_result.log')
if save_settings:
# ------------------------------- #
# Saving training parameters
# ------------------------------- #
result_logger.info(f'Model: {model_arch}')
result_logger.info(f'Layer: {args.layer}')
result_logger.info(f'Learning rate: {learning_rate}')
result_logger.info(f'alpha: {alpha}')
result_logger.info(f'Normalize feature vector: {normalize}')
# ------------------------------- #
# Load extracted knowledge graph
# ------------------------------- #
knowledge_graph = Graph()
classFile_to_superclasses, superclassID_to_wikiID =\
knowledge_graph.class_file_to_superclasses(1, [1,2])
####################
# Prepare Data Set #
####################
print('preparing dataset')
base_cls, val_cls, support_cls = get_splits()
base = MiniImageNet('base', base_cls, val_cls, support_cls,
classFile_to_superclasses)
base_loader = DataLoader(base, batch_size=256, shuffle=True, num_workers=4)
support = MiniImageNet('support',
base_cls,
val_cls,
support_cls,
classFile_to_superclasses,
eval=True)
support_loader_1 = DataLoader(support,
batch_sampler=SupportingSetSampler(
support, 1, 5, 15, support_groups),
num_workers=4)
support_loader_5 = DataLoader(support,
batch_sampler=SupportingSetSampler(
support, 5, 5, 15, support_groups),
num_workers=4)
#########
# Model #
#########
if model_arch == 'conv4':
model = models.Conv4Attension(len(base_cls),
len(superclassID_to_wikiID))
if model_arch == 'resnet10':
model = models.resnet10(len(base_cls), len(superclassID_to_wikiID))
if model_arch == 'resnet18':
model = models.resnet18(len(base_cls), len(superclassID_to_wikiID))
model.to(device)
# loss function and optimizer
criterion = loss_fn(alpha)
optimizer = torch.optim.SGD(model.parameters(),
lr=learning_rate,
momentum=0.9,
weight_decay=1e-4,
nesterov=True)
scheduler = MultiStepLR(optimizer,
milestones=scheduler_milestones,
gamma=0.1)
if save_settings:
result_logger.info(
'optimizer: torch.optim.SGD(model.parameters(), '
f'lr={learning_rate}, momentum=0.9, weight_decay=1e-4, nesterov=True)'
)
result_logger.info(
f'scheduler: MultiStepLR(optimizer, milestones={scheduler_milestones}, gamma=0.1)\n'
)
# result_logger.info('='*40+'Results Below'+'='*40+'\n')
if checkpoint:
print('load model...')
model.load_state_dict(torch.load(f'{model_dir}/{start_epoch-1}.pth'))
model.to(device)
for _ in range(start_epoch - 1):
scheduler.step()
# ------------------------------- #
# Start to train
# ------------------------------- #
if toTrain:
for epoch in range(start_epoch, start_epoch + num_epoch):
model.train()
train(model, normalize, base_loader, optimizer, criterion, epoch,
start_epoch + num_epoch - 1, device, train_logger)
scheduler.step()
if epoch % model_saving_rate == 0:
torch.save(model.state_dict(), f'{model_dir}/{epoch}.pth')
# ------------------------------- #
# Evaluate current model
# ------------------------------- #
if toEvaluate:
if epoch % evaluation_rate == 0:
evaluate(model, normalize, epoch, support_loader_1, 1, 5,
15, device, result_logger)
evaluate(model, normalize, epoch, support_loader_5, 5, 5,
15, device, result_logger)
else:
if toEvaluate:
evaluate(model, normalize, start_epoch - 1, support_loader_1, 1, 5,
15, device, result_logger)
evaluate(model, normalize, start_epoch - 1, support_loader_5, 5, 5,
15, device, result_logger)
result_logger.info('=' * 140)
def main(args):
if args.checkpoint == '':
args.checkpoint = "checkpoints/ic15_%s_bs_%d_ep_%d" % (
args.arch, args.batch_size, args.n_epoch)
if args.pretrain:
if 'synth' in args.pretrain:
args.checkpoint += "_pretrain_synth"
else:
args.checkpoint += "_pretrain_ic17"
print('checkpoint path: %s' % args.checkpoint)
print('init lr: %.8f' % args.lr)
print('schedule: ', args.schedule)
sys.stdout.flush()
if not os.path.isdir(args.checkpoint):
os.makedirs(args.checkpoint)
kernel_num = 1
min_scale = 0.4
start_epoch = 0
data_loader = IC15Loader(root_dir=args.root_dir,
is_transform=True,
img_size=args.img_size,
kernel_num=kernel_num,
min_scale=min_scale)
train_loader = torch.utils.data.DataLoader(data_loader,
batch_size=args.batch_size,
shuffle=True,
num_workers=3,
drop_last=True,
pin_memory=True)
if args.arch == "resnet50":
model = models.resnet50(pretrained=True, num_classes=kernel_num)
elif args.arch == "resnet101":
model = models.resnet101(pretrained=True, num_classes=kernel_num)
elif args.arch == "resnet152":
model = models.resnet152(pretrained=True, num_classes=kernel_num)
elif args.arch == "resnet18":
model = models.resnet18(pretrained=True, num_classes=kernel_num)
model = torch.nn.DataParallel(model).cuda()
if hasattr(model.module, 'optimizer'):
optimizer = model.module.optimizer
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.99,
weight_decay=5e-4)
title = 'icdar2015'
if args.pretrain:
print('Using pretrained model.')
assert os.path.isfile(
args.pretrain), 'Error: no checkpoint directory found!'
checkpoint = torch.load(args.pretrain)
model.load_state_dict(checkpoint['state_dict'])
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Learning Rate', 'Train Loss', 'Train Acc.', 'Train IOU.'])
elif args.resume:
print('Resuming from checkpoint.')
assert os.path.isfile(
args.resume), 'Error: no checkpoint directory found!'
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print('Training from scratch.')
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Learning Rate', 'Train Loss', 'Train Acc.', 'Train IOU.'])
writer = SummaryWriter(logdir=args.checkpoint)
for epoch in range(start_epoch, args.n_epoch):
adjust_learning_rate(args, optimizer, epoch)
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.n_epoch, optimizer.param_groups[0]['lr']))
train_loss, train_te_acc, train_te_iou = train(train_loader, model,
dice_loss, optimizer,
epoch, writer)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'lr': args.lr,
'optimizer': optimizer.state_dict(),
},
checkpoint=args.checkpoint)