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))