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
from config import cfg
import torch
import torch.nn as nn
import numpy as np
import torch.optim as optim
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