def train_net(args):
torch.manual_seed(7)
np.random.seed(7)
checkpoint = args.checkpoint
start_epoch = 0
best_acc = 0
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':
model = MobileNet(1.0)
elif args.network == 'mr18':
print("mr18")
model = myResnet18()
else:
model = resnet_face18(args.use_se)
model = nn.DataParallel(model)
metric_fc = ArcMarginModel(args)
metric_fc = nn.DataParallel(metric_fc)
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']
logger = get_logger()
# 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)
scheduler = StepLR(optimizer, step_size=args.lr_step, gamma=0.1)
# Epochs
for epoch in range(start_epoch, args.end_epoch):
scheduler.step()
if args.full_log:
lfw_acc, threshold = lfw_test(model)
writer.add_scalar('LFW_Accuracy', lfw_acc, epoch)
full_log(epoch)
start = datetime.now()
# One epoch's training
train_loss, train_top5_accs = train(train_loader=train_loader,
model=model,
metric_fc=metric_fc,
criterion=criterion,
optimizer=optimizer,
epoch=epoch,
logger=logger,
writer=writer)
writer.add_scalar('Train_Loss', train_loss, epoch)
writer.add_scalar('Train_Top5_Accuracy', train_top5_accs, epoch)
end = datetime.now()
delta = end - start
print('{} seconds'.format(delta.seconds))
# One epoch's validation
lfw_acc, threshold = lfw_test(model)
writer.add_scalar('LFW Accuracy', lfw_acc, epoch)
# Check if there was an improvement
is_best = lfw_acc > best_acc
best_acc = max(lfw_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("--mode", type=str, default="train")
parser.add_argument("--model", type=str, default="mobilenet_v2")
parser.add_argument("--dataset", type=str, default="cifar10")
parser.add_argument("--dataroot", type=str, default="/tmp/data")
parser.add_argument("--batch_size", type=int, default=128)
parser.add_argument("--n_epochs", type=int, default=100)
parser.add_argument("--lr", type=float, default=1e-3)
parser.add_argument("--n_gpus", type=int, default=1)
parser.add_argument("--checkpoint", type=str, default="/tmp/chkpt.pth.tar")
parser.add_argument("--save_every", type=int, default=10)
parser.add_argument("--pretrained", type=str, default=None)
args = parser.parse_args()
print(args)
if torch.cuda.is_available():
print("cuda is available, use cuda")
device = torch.device("cuda")
else:
print("cuda is not available, use cpu")
device = torch.device("cpu")
print("download dataset: {}".format(args.dataset))
train_loader, test_loader, n_classes = get_loaders(
dataset=args.dataset, root=args.dataroot, batch_size=args.batch_size)
print("build model: {}".format(args.model))
if args.model == "mobilenet":
from models import MobileNet
model = MobileNet(n_classes=n_classes)
elif args.model == "mobilenet_v2":
from models import MobileNet_v2
model = MobileNet_v2(n_classes=n_classes)
elif args.model == "shufflenet":
from models import ShuffleNet
model = ShuffleNet(n_classes=n_classes)
elif args.model == "shufflenet_v2":
from models import ShuffleNet_v2
model = ShuffleNet_v2(n_classes=n_classes)
elif args.model == "squeezenet":
from models import SqueezeNet
model = SqueezeNet(n_classes=n_classes)
else:
raise NotImplementedError
model = model.to(device)
if args.pretrained:
model.load_state_dict(torch.load(args.checkpoint))
if args.n_gpus > 1:
gpus = []
for i in range(args.n_gpus):
gpus.append(i)
model = nn.DataParallel(model, device_ids=gpus)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
criterion = nn.CrossEntropyLoss()
if args.mode == "train":
for epoch in range(args.n_epochs):
train(epoch, model, optimizer, criterion, train_loader, device)
if (epoch + 1) % args.save_every == 0:
print("saving model...")
torch.save(the_model.state_dict(), args.checkpoint)
elif args.mode == "test":
test(model, criterion, test_loader, device)
else:
raise NotImplementedError
def train_net(args):
torch.manual_seed(7)
np.random.seed(7)
checkpoint = args.checkpoint
start_epoch = 0
best_acc = 0
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':
model = MobileNet(1.0)
else:
model = resnet_face18(args.use_se)
model = nn.DataParallel(model)
metric_fc = ArcMarginModel(args)
metric_fc = nn.DataParallel(metric_fc)
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)
optimizer = InsightFaceOptimizer(
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']
logger = get_logger()
# 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=8)
# Epochs
for epoch in range(start_epoch, args.end_epoch):
# 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,
logger=logger)
print('\nCurrent effective learning rate: {}\n'.format(optimizer.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', optimizer.lr, epoch)
# 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 train_net(args):
torch.manual_seed(7) #torch的随机种子，在torch.randn使用
np.random.seed(7)
checkpoint = args.checkpoint
start_epoch = 0
best_acc = 0
writer = SummaryWriter() #tensorboard
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':
model = MobileNet(1.0)
else:
model = resnet_face18(args.use_se)
model = nn.DataParallel(model)
metric_fc = ArcMarginModel(args)
metric_fc = nn.DataParallel(metric_fc)
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)
optimizer = InsightFaceOptimizer(
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']
logger = get_logger()
# 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 = Dataset(root=args.train_path,
phase='train',
input_shape=(3, 112, 112))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=8)
# Epochs
for epoch in range(start_epoch, args.end_epoch):
# 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,
logger=logger)
print('\nCurrent effective learning rate: {}\n'.format(optimizer.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', optimizer.lr, epoch)
# Save checkpoint
if epoch % 10 == 0:
save_checkpoint(epoch, epochs_since_improvement, model, metric_fc,
optimizer, best_acc)
def main():
parser = argparse.ArgumentParser(
description='Deep Learning SNN simulation')
parser.add_argument('--config-file', default='config_1.ini')
args = parser.parse_args()
print(args)
config = configparser.ConfigParser()
config.read(args.config_file)
pprint.pprint(
{section: dict(config[section])
for section in config.sections()})
print
defaults = config['DEFAULT']
device = defaults['device']
app_name = defaults['app_name']
print('[INFO] Simulating spiking {}'.format(app_name))
if not os.path.isdir(app_name):
os.mkdir(app_name)
out_dir = app_name
org_model = config['original model']
" Load the original model "
net = None
if 'mobilenet_cifar10' in org_model['arch']:
net = MobileNet()
if org_model['arch'] == 'vgg_cifar10':
net = VGG_16_cifar10
num_classes = org_model.getint('num_classes')
model_path = org_model['model_path']
file_name = org_model['file_name']
state = None
state, net = load_model(net, model_path, file_name)
net = net.to(device)
" Load the dataset "
testloader, img_size = None, None
data_config = config['dataset']
data_dir = data_config['data_dir']
if data_config['dataset'] == 'cifar10':
_, testloader = load_cifar10(data_dir, org_model['arch'])
img_size = [-1, 3, 32, 32]
" Tasks to do "
tasks = config['functions']
if tasks.getboolean('validate'):
print(net)
validate(net, testloader, device)
new_net = None
if tasks.getboolean('remove_bn'):
if has_bn(net):
new_net = MobileNet_nobn()
new_net = merge_bn(net, new_net)
new_net = new_net.to(device)
print(new_net)
print('Validating model after folding back BN layers...')
validate(new_net, testloader, device)
save_model(new_net, state, out_dir, 'nobn_' + file_name)
else:
print('model has no BN layers')
if tasks.getboolean('use_nobn'):
if 'mobilenet_cifar10' in org_model['arch']:
net = MobileNet_nobn()
state, net = load_model(net, out_dir, 'nobn_' + file_name)
net = net.to(device)
" quantize model "
quant_config = config['quantization']
W = quant_config.getint('W')
F = quant_config.getint('F')
if tasks.getboolean('quantize'):
net = quantize_model(net, W, F)
net = net.to(device)
print('Validating model after quantization...')
state['acc'] = validate(net, testloader, device)
save_model(net, state, out_dir, 'quant{}.{}_'.format(W, F) + file_name)
" use quantized model "
if tasks.getboolean('use_quant'):
state, net = load_model(net, out_dir,
'quant{}.{}_'.format(W, F) + file_name)
net = net.to(device)
if tasks.getboolean('validate_nobn'):
if not has_bn(net):
print('Validating no_bn model...')
validate(net, testloader)
else:
print('model has BN layers!! Exiting..')
exit()
" compute thresholds "
spike_config = config['spiking']
percentile = spike_config.getfloat('percentile')
if spike_config.getboolean('compute_thresholds'):
compute_thresholds(net, testloader, out_dir, percentile, device)
" convert ann to snn "
spike_net, thresholds, max_acts, clamp_slope = None, None, None, None
if spike_config.getboolean('convert_to_spike'):
thresholds = np.loadtxt(os.path.join(out_dir, 'thresholds.txt'))
max_acts = np.loadtxt(os.path.join(out_dir, 'max_acts.txt'))
clamp_slope = spike_config.getfloat('clamp_slope')
spike_net = createSpikingModel(net, org_model['arch'], num_classes,
spike_config, thresholds, max_acts,
device)
print(spike_net)
sanity_check(net, spike_net, max_acts)
" simulate snn "
if spike_config.getboolean('simulate_spiking'):
thresholds = np.loadtxt(os.path.join(out_dir, 'thresholds.txt'))
max_acts = np.loadtxt(os.path.join(out_dir, 'max_acts.txt'))
simulate_spike_model(net, org_model['arch'], testloader, config,
thresholds, max_acts, num_classes, img_size,
device)
" plot correlations "
if spike_config.getboolean('plot_correlations'):
corr = np.load(os.path.join(out_dir, 'layerwise_corr.npy'))
plot_config = config['plotting']
#print('corr matrix shape: {}'.format(corr.shape))
plot_correlations(corr, out_dir, plot_config)
])
# Argv
test_fpath = sys.argv[1]
model_fpath = sys.argv[2]
output_fpath = sys.argv[3]
print('# [Info] Argv')
print(' - Test : {}'.format(test_fpath))
print(' - Model : {}'.format(model_fpath))
print(' = Output : {}'.format(output_fpath))
# Make data loader
test_dataset = MyDataset(filename=test_fpath, is_train=False, transform=test_transform)
test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=False)
# Load model
model = MobileNet()
model.load_state_dict(torch.load(model_fpath, map_location=device))
model.to(device)
# Model prediction
model.eval()
prediction = []
for i, data in enumerate(test_loader):
data_device = data.to(device)
output = model(data_device)
labels = torch.max(output, 1)[1]
for label in labels:
prediction.append(label)
# Output prediction
print('# [Info] Output prediction: {}'.format(output_fpath))
with open(output_fpath, 'w') as f:
f.write('id,label\n')
for i, v in enumerate(prediction):
f.write('%d,%d\n' %(i, v))