def get_pack():
set_seeds()
pack = recover_pack()
model_dict = torch.load(
'./logs/vgg16_cifar10/ckp.160.torch',
map_location='cpu' if not cfg.base.cuda else 'cuda')
pack.net.module.load_state_dict(model_dict)
GBNs = GatedBatchNorm2d.transform(pack.net)
for gbn in GBNs:
gbn.extract_from_bn()
pack.optimizer = optim.SGD(pack.net.parameters(),
lr=2e-3,
momentum=cfg.train.momentum,
weight_decay=cfg.train.weight_decay,
nesterov=cfg.train.nesterov)
return pack, GBNs
def clone_model(net):
model = get_model()
gbns = GatedBatchNorm2d.transform(model.module)
model.load_state_dict(net.state_dict())
return model, gbns
from logger import logger
from main import set_seeds, recover_pack, _step_lr, _sgdr
from models import get_model
from utils import dotdict
from prune.universal import Meltable, GatedBatchNorm2d, Conv2dObserver, IterRecoverFramework, FinalLinearObserver
from prune.utils import analyse_model, finetune
set_seeds()
pack = recover_pack()
# model_dict = torch.load('logs/resnet56m_cifar100_baseline_{}4/ckp.398.torch'.format(res), map_location='cpu' if not cfg.base.cuda else 'cuda')
# pack.net.module.load_state_dict(model_dict)
GBNs = GatedBatchNorm2d.transform(pack.net)
for gbn in GBNs:
gbn.extract_from_bn()
pack.optimizer = optim.SGD(
pack.net.parameters() ,
lr=2e-3,
momentum=cfg.train.momentum,
weight_decay=cfg.train.weight_decay,
nesterov=cfg.train.nesterov
)
print(cfg.base.task_name)
import uuid
def bottleneck_set_group(net):
layers = [