num_workers=args.num_workers) # 测试集数据 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') if args.refine: print('******Refine model******') #checkpoint = torch.load('../prune/models_save/nin_refine.pth') checkpoint = torch.load(args.refine) if args.model_type == 0: model = nin.Net(cfg=checkpoint['cfg'], a_bits=args.a_bits, w_bits=args.w_bits) else: model = nin_gc.Net(cfg=checkpoint['cfg'], a_bits=args.a_bits, w_bits=args.w_bits) model.load_state_dict(checkpoint['state_dict']) best_acc = 0 else: print('******Initializing model******') if args.model_type == 0: model = nin.Net(a_bits=args.a_bits, w_bits=args.w_bits) else: model = nin_gc.Net(a_bits=args.a_bits, w_bits=args.w_bits) best_acc = 0 for m in model.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_uniform_(m.weight.data) m.bias.data.zero_() elif isinstance(m, nn.Linear):
trainset = torchvision.datasets.CIFAR10(root = args.data, train = True, download = True, transform = transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.num_workers) # 训练集数据 testset = torchvision.datasets.CIFAR10(root = args.data, train = False, download = True, transform = transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=args.eval_batch_size, shuffle=False, num_workers=args.num_workers) # 测试集数据 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') if args.refine: print('******Refine model******') #checkpoint = torch.load('../prune/models_save/nin_refine.pth') checkpoint = torch.load(args.refine) if args.model_type == 0: model = nin.Net(cfg=checkpoint['cfg'], abits=args.Abits, wbits=args.Wbits) else: model = nin_gc.Net(cfg=checkpoint['cfg'], abits=args.Abits, wbits=args.Wbits) model.load_state_dict(checkpoint['state_dict']) best_acc = 0 else: print('******Initializing model******') if args.model_type == 0: model = nin.Net(abits=args.Abits, wbits=args.Wbits) else: model = nin_gc.Net(abits=args.Abits, wbits=args.Wbits) best_acc = 0 for m in model.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_uniform_(m.weight.data) m.bias.data.zero_() elif isinstance(m, nn.Linear): m.weight.data.normal_(0, 0.01)
transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=args.eval_batch_size, shuffle=False, num_workers=args.num_workers) classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') if args.prune_quant: print('******Prune Quant model******') #checkpoint = torch.load('../prune/models_save/nin_refine.pth') checkpoint = torch.load(args.prune_quant) cfg = checkpoint['cfg'] if args.model_type == 0: model = nin.Net(cfg=checkpoint['cfg']) else: model = nin_gc.Net(cfg=checkpoint['cfg']) model.load_state_dict(checkpoint['state_dict']) best_acc = 0 print('***ori_model***\n', model) quantize.prepare(model, inplace=True, a_bits=args.a_bits, w_bits=args.w_bits, q_type=args.q_type, q_level=args.q_level, device=device, weight_observer=args.weight_observer, bn_fuse=args.bn_fuse, bn_fuse_cali=args.bn_fuse_cali, pretrained_model=args.pretrained_model, qaft=args.qaft, ptq=args.ptq, percentile=args.percentile) print('\n***quant_model***\n', model) elif args.prune_qaft:
parser.add_argument('--A', type=int, default=2, help='Ab:2, Afp:32') args = parser.parse_args() print('==> Options:', args) if args.gpu_id: os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id if args.prune_quant: print('******Prune Quant model******') if args.model_type == 0: checkpoint = torch.load('../models_save/nin.pth') quant_model_train = nin.Net(cfg=checkpoint['cfg']) else: checkpoint = torch.load('../models_save/nin_gc.pth') quant_model_train = nin_gc.Net(cfg=checkpoint['cfg']) else: if args.model_type == 0: checkpoint = torch.load('../models_save/nin.pth') quant_model_train = nin.Net() else: checkpoint = torch.load('../models_save/nin_gc.pth') quant_model_train = nin_gc.Net() quant_bn_fused_model_inference = copy.deepcopy(quant_model_train) quantize.prepare(quant_model_train, inplace=True, A=args.A, W=args.W) quantize.prepare(quant_bn_fused_model_inference, inplace=True, A=args.A, W=args.W, quant_inference=True) quant_model_train.load_state_dict(checkpoint['state_dict'])
trainset = torchvision.datasets.CIFAR10(root = args.data, train = True, download = True, transform = transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.num_workers) # 训练集数据 testset = torchvision.datasets.CIFAR10(root = args.data, train = False, download = True, transform = transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=args.eval_batch_size, shuffle=False, num_workers=args.num_workers) # 测试集数据 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') if args.refine: print('******Refine model******') #checkpoint = torch.load('../prune/models_save/nin_refine.pth') checkpoint = torch.load(args.refine) if args.model_type == 0: model = nin.Net(cfg=checkpoint['cfg'], abits=args.Abits, wbits=args.Wbits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level) else: model = nin_gc.Net(cfg=checkpoint['cfg'], abits=args.Abits, wbits=args.Wbits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level) model_dict = model.state_dict() update_state_dict = {k:v for k,v in checkpoint['state_dict'].items() if k in model_dict.keys()} model_dict.update(update_state_dict) print('fp32_model weight load successfully') model.load_state_dict(model_dict) best_acc = 0 else: print('******Initializing model******') if args.model_type == 0: model = nin.Net(abits=args.Abits, wbits=args.Wbits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level) else: model = nin_gc.Net(abits=args.Abits, wbits=args.Wbits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level) best_acc = 0 for m in model.modules(): if isinstance(m, nn.Conv2d):
help='disables CUDA training') parser.add_argument('--percent', type=float, default=0.4, help='nin_gc:0.4') parser.add_argument('--layers', type=int, default=9, help='layers (default: 9)') parser.add_argument('--model', default='models_save/nin_gc_preprune.pth', type=str, metavar='PATH', help='path to raw trained model (default: none)') args = parser.parse_args() layers = args.layers print(args) model = nin_gc.Net() if args.model: if os.path.isfile(args.model): print("=> loading checkpoint '{}'".format(args.model)) model.load_state_dict(torch.load(args.model)['state_dict']) else: print("=> no checkpoint found at '{}'".format(args.resume)) print('旧模型: ', model) total = 0 i = 0 for m in model.modules(): if isinstance(m, nn.BatchNorm2d): if i < layers - 1: i += 1 total += m.weight.data.shape[0]
shuffle=False, num_workers=args.num_workers) # 测试集数据 # cifar10类别 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') # model if args.refine: print('******Refine model******') #checkpoint = torch.load('../prune/models_save/nin_refine.pth') checkpoint = torch.load(args.refine) if args.model_type == 0: model = nin.Net(cfg=checkpoint['cfg'], A=args.A, W=args.W) else: model = nin_gc.Net(cfg=checkpoint['cfg'], A=args.A, W=args.W) model.load_state_dict(checkpoint['state_dict']) best_acc = 0 else: print('******Initializing model******') # ******************** 在model的量化卷积中同时量化A(特征)和W(模型参数) ************************ if args.model_type == 0: model = nin.Net(A=args.A, W=args.W) else: model = nin_gc.Net(A=args.A, W=args.W) #model = nin_bn_conv.Net() best_acc = 0 for m in model.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_uniform_(m.weight.data) m.bias.data.zero_()
#checkpoint = torch.load('../prune/models_save/nin_refine.pth') checkpoint = torch.load(args.refine) if args.model_type == 0: model = nin.Net(cfg=checkpoint['cfg'], a_bits=args.a_bits, w_bits=args.w_bits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level, device=device, weight_observer=args.weight_observer) else: model = nin_gc.Net(cfg=checkpoint['cfg'], a_bits=args.a_bits, w_bits=args.w_bits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level, device=device, weight_observer=args.weight_observer) model_dict = model.state_dict() update_state_dict = { k: v for k, v in checkpoint['state_dict'].items() if k in model_dict.keys() } model_dict.update(update_state_dict) print('fp32_model weight load successfully') model.load_state_dict(model_dict) best_acc = 0 else: print('******Initializing model******')
'ship', 'truck') if args.refine: print('******Refine model******') #checkpoint = torch.load('../prune/models_save/nin_refine.pth') checkpoint = torch.load(args.refine) if args.model_type == 0: model = nin.Net(cfg=checkpoint['cfg'], abits=args.Abits, wbits=args.Wbits, bn_fold=args.bn_fold, q_type=args.q_type) else: model = nin_gc.Net(cfg=checkpoint['cfg'], abits=args.Abits, wbits=args.Wbits, bn_fold=args.bn_fold, q_type=args.q_type) model.load_state_dict(checkpoint['state_dict']) best_acc = 0 else: print('******Initializing model******') if args.model_type == 0: model = nin.Net(abits=args.Abits, wbits=args.Wbits, bn_fold=args.bn_fold, q_type=args.q_type) else: model = nin_gc.Net(abits=args.Abits, wbits=args.Wbits, bn_fold=args.bn_fold,
trainset = torchvision.datasets.CIFAR10(root = args.data, train = True, download = True, transform = transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.num_workers) # 训练集数据 testset = torchvision.datasets.CIFAR10(root = args.data, train = False, download = True, transform = transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=args.eval_batch_size, shuffle=False, num_workers=args.num_workers) # 测试集数据 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') if args.refine: print('******Refine model******') #checkpoint = torch.load('../prune/models_save/nin_refine.pth') checkpoint = torch.load(args.refine) if args.model_type == 0: model = nin.Net(cfg=checkpoint['cfg'], a_bits=args.a_bits, w_bits=args.w_bits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level) else: model = nin_gc.Net(cfg=checkpoint['cfg'], a_bits=args.a_bits, w_bits=args.w_bits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level) model_dict = model.state_dict() update_state_dict = {k:v for k,v in checkpoint['state_dict'].items() if k in model_dict.keys()} model_dict.update(update_state_dict) print('fp32_model weight load successfully') model.load_state_dict(model_dict) best_acc = 0 else: print('******Initializing model******') if args.model_type == 0: model = nin.Net(a_bits=args.a_bits, w_bits=args.w_bits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level) else: model = nin_gc.Net(a_bits=args.a_bits, w_bits=args.w_bits, bn_fuse=args.bn_fuse, q_type=args.q_type, q_level=args.q_level) best_acc = 0 for m in model.modules(): if isinstance(m, nn.Conv2d):
if args.prune_refine: print('******Prune Refine model******') #checkpoint = torch.load('models_save/nin_prune.pth') checkpoint = torch.load(args.prune_refine) cfg = checkpoint['cfg'] model = nin.Net(cfg=checkpoint['cfg']) model.load_state_dict(checkpoint['state_dict']) best_acc = 0 elif args.refine: print('******Float Refine model******') #checkpoint = torch.load('models_save/nin.pth') state_dict = torch.load(args.refine) if args.model_type == 0: model = nin.Net() else: model = nin_gc.Net() model.load_state_dict(state_dict) best_acc = 0 elif args.resume: print('******Reume model******') #checkpoint = torch.load('models_save/nin.pth') #checkpoint = torch.load('models_save/nin_sparse.pth') checkpoint = torch.load(args.resume) if args.model_type == 0: model = nin.Net() else: model = nin_gc.Net() model.load_state_dict(checkpoint['state_dict']) best_acc = checkpoint['best_acc'] else: # nin_gc_retrain
'ship', 'truck') if args.refine: print('******Refine model******') #checkpoint = torch.load('models_save/nin_prune.pth') checkpoint = torch.load(args.refine) cfg = checkpoint['cfg'] model = nin.Net(cfg=checkpoint['cfg'], quant_type=args.quant_type) model.load_state_dict(checkpoint['state_dict']) best_acc = 0 else: # nin_gc_retrain if args.gc_refine: print('******Refine model******') cfg = args.gc_refine model = nin_gc.Net(cfg=cfg, quant_type=args.quant_type) else: print('******Initializing model******') if args.model_type == 0: model = nin.Net(quant_type=args.quant_type) else: model = nin_gc.Net(quant_type=args.quant_type) best_acc = 0 for m in model.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_uniform_(m.weight.data) m.bias.data.zero_() elif isinstance(m, nn.Linear): m.weight.data.normal_(0, 0.01) m.bias.data.zero_()
if args.refine: print('******Refine model******') #checkpoint = torch.load('models_save/nin_prune.pth') checkpoint = torch.load(args.refine) cfg = checkpoint['cfg'] model = nin.Net(cfg=checkpoint['cfg']) model.load_state_dict(checkpoint['state_dict']) best_acc = 0 else: print('******Initializing model******') model = nin.Net() if args.type == 0: model = nin.Net() elif args.type == 1: model = nin_gc.Net() elif args.type == 2: model = standard_dw.Net() ''' cfg = [] #gc_prune —— cfg model = nin_gc.Net(cfg=cfg) ''' best_acc = 0 for m in model.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_uniform_(m.weight.data) m.bias.data.zero_() elif isinstance(m, nn.Linear): m.weight.data.normal_(0, 0.01) m.bias.data.zero_() if args.resume: