from csgd.csgd_pipeline import csgd_iterative
from constants import VGG_ORIGIN_DEPS, LRSchedule
from base_config import get_baseconfig_by_epoch
from constants import unet_succeeding_strategy, unet_pacesetter_dict, unet_origin_deps_flattened
from network.UNet import UNet
from util import loadYaml, parseArgs
from network.UNet import UNet
base_lrs = LRSchedule(base_lr=5e-2, max_epochs=600, lr_epoch_boundaries=[200, 400], lr_decay_factor=0.1, linear_final_lr=None)
csgd_lrs = LRSchedule(base_lr=3e-2, max_epochs=600, lr_epoch_boundaries=[200, 400], lr_decay_factor=0.1, linear_final_lr=None)
# csgd_lrs = LRSchedule(base_lr=1e-2, max_epochs=600, lr_epoch_boundaries=[200, 400], lr_decay_factor=0.1, linear_final_lr=None)
args = parseArgs()
unetcfg, saveName = loadYaml(args.config)
def csgd_unet():
try_arg = 'slim_5-8'
network_type = 'unet'
# dataset_name = 'cifar10'
dataset_name = 'GoPro'
base_log_dir = 'csgd_exps/{}_{}_base'.format(network_type, try_arg)
csgd_log_dir = 'csgd_exps/{}_{}_csgd'.format(network_type, try_arg)
weight_decay_strength = 1e-4
batch_size = 64
unet = UNet(3, 3)
origin_deps = unet_origin_deps_flattened(unet)
centri_strength = 3e-3
# 3 phase iter to slim
start_arg = parser.parse_args()
network_type = start_arg.arch
block_type = start_arg.block_type
conti_or_fs = start_arg.conti_or_fs
assert conti_or_fs in ['continue', 'fs']
assert block_type in ['acb', 'base']
auto_continue = conti_or_fs == 'continue'
print('auto continue: ', auto_continue)
gamma_init = None
if network_type == 'sres18':
weight_decay_strength = 1e-4
batch_size = 256
lrs = LRSchedule(base_lr=0.1, max_epochs=100, lr_epoch_boundaries=None, lr_decay_factor=None,
linear_final_lr=None, cosine_minimum=0)
warmup_epochs = 0
gamma_init = 1
elif network_type == 'sres34':
weight_decay_strength = 1e-4
batch_size = 256
lrs = LRSchedule(base_lr=0.1, max_epochs=100, lr_epoch_boundaries=None, lr_decay_factor=None,
linear_final_lr=None, cosine_minimum=0)
warmup_epochs = 0
gamma_init = 1
elif network_type == 'sres50':
weight_decay_strength = 1e-4
batch_size = 256
lrs = LRSchedule(base_lr=0.1, max_epochs=100, lr_epoch_boundaries=None, lr_decay_factor=None,