def test():
#model = torchvision.models.alexnet(pretrained=True)
#input_shape = (3, 224, 224)
#model = torchvision.models.lenet(pretrained=True)
cfg = [300, 100]
model = models.LeNet_300_100(bias_flag=True, cfg=cfg)
input_shape = [784]
new_model = neuron_merge(model, input_shape, threshold=0.5)
print(model)
print(new_model)
datasets.MNIST('data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=args.test_batch_size, shuffle=True, **kwargs)
# generate the model
if args.arch == 'LeNet_300_100':
model = models.LeNet_300_100(args.prune)
elif args.arch == 'LeNet_5':
model = models.LeNet_5(args.prune)
else:
print('ERROR: specified arch is not suppported')
exit()
if not args.pretrained:
best_acc = 0.0
else:
pretrained_model = torch.load(args.pretrained)
best_acc = pretrained_model['acc']
load_state(model, pretrained_model['state_dict'])
if args.cuda:
model.cuda()
cfg[i] = int(cfg[i] * (1 - args.pruning_ratio))
temp_cfg[i] = cfg[i] * args.depth_wide[1]
elif args.target == 'ip':
if args.arch == 'LeNet_300_100':
cfg = [300, 100]
for i in range(len(cfg)):
cfg[i] = round(cfg[i] * (1 - args.pruning_ratio))
temp_cfg = cfg
pass
# generate the model
if args.arch == 'VGG':
model = models.VGG(num_classes, cfg=cfg)
elif args.arch == 'LeNet_300_100':
model = models.LeNet_300_100(bias_flag=True, cfg=cfg)
elif args.arch == 'ResNet':
model = models.ResNet(int(args.depth_wide), num_classes, cfg=cfg)
elif args.arch == 'WideResNet':
model = models.WideResNet(args.depth_wide[0],
num_classes,
widen_factor=args.depth_wide[1],
cfg=cfg)
else:
pass
if args.cuda:
model.cuda()
# pretrain
best_acc = 0.0
parser.add_argument('--specialize', type=int, metavar='N', default=None,
help='extract the specialized network')
parser.add_argument('--algo',default=None,
help='clustering algorithm to use')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
print_args(args)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
# generate the model
if args.arch == 'LeNet_300_100':
model = models.LeNet_300_100()
elif args.arch == 'LeNet_5':
model = models.LeNet_5()
else:
print('ERROR: specified arch is not suppported')
exit()
if not args.pretrained:
best_acc = 0.0
else:
pretrained_model = torch.load(args.pretrained)
# best_acc = pretrained_model['acc'] if not args.algo else 0
best_acc = 0.0
mask = pretrained_model['mask'] if 'mask' in pretrained_model else []
model.mask = mask
load_state(model, pretrained_model['state_dict'])
