def loadModel(self, model_dir=constant.train_config['trained_model_dir'], model_name=constant.predict_config['best_model_name']): model_path = model_dir + os.sep + model_name classifier_model = Model(len(self.label_encoder.classes_)) nn_model = classifier_model.createModel() nn_model.load_state_dict(torch.load(model_path)) nn_model.eval() use_cuda = torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") return nn_model.to(device)
def __init__(self, batch_size=constant.train_config['train_batch_size'], num_epochs=constant.train_config['num_epochs'], dataset_path=constant.dataset_config['dataset_file'], features=constant.dataset_config['features'], labels=constant.dataset_config['labels']): #traing parameters self.batch_size = batch_size self.num_epochs = num_epochs #data-loader for training in pytorch dataloader format self.train_batch, self.valid_batch, self.target_num_classes = self.loadData(dataset_path, features, labels) #NN Model functions self.classifier_model = Model(self.target_num_classes) self.error, self.optimizer = self.classifier_model.error_optimizer() #logging description self.model_name = constant.model_config['model_name']
def sample(cfg, logger): model = Model(startf=cfg.MODEL.START_CHANNEL_COUNT, layer_count=cfg.MODEL.LAYER_COUNT, maxf=cfg.MODEL.MAX_CHANNEL_COUNT, latent_size=cfg.MODEL.LATENT_SPACE_SIZE, truncation_psi=cfg.MODEL.TRUNCATIOM_PSI, truncation_cutoff=cfg.MODEL.TRUNCATIOM_CUTOFF, mapping_layers=cfg.MODEL.MAPPING_LAYERS, channels=3) model.eval() logger.info("Trainable parameters generator:") count_parameters(model.generator) model_dict = { 'generator_s': model.generator, 'mapping_fl_s': model.mapping, 'dlatent_avg': model.dlatent_avg, } checkpointer = Checkpointer(cfg, model_dict, logger=logger, save=True) checkpointer.load() ctx = bimpy.Context() remove = bimpy.Bool(False) layers = bimpy.Int(8) ctx.init(1800, 1600, "Styles") rnd = np.random.RandomState(5) latents = rnd.randn(1, cfg.MODEL.LATENT_SPACE_SIZE) sample = torch.tensor(latents).float().cuda() def update_image(sample): with torch.no_grad(): torch.manual_seed(0) model.eval() x_rec = model.generate(layers.value, remove.value, z=sample) #model.generator.set(l.value, c.value) resultsample = ((x_rec * 0.5 + 0.5) * 255).type(torch.long).clamp( 0, 255) resultsample = resultsample.cpu()[0, :, :, :] return resultsample.type(torch.uint8).transpose(0, 2).transpose(0, 1) with torch.no_grad(): save_image(model.generate(8, True, z=sample) * 0.5 + 0.5, 'sample.png') im = bimpy.Image(update_image(sample)) while (not ctx.should_close()): with ctx: bimpy.set_window_font_scale(2.0) if bimpy.checkbox('REMOVE BLOB', remove): im = bimpy.Image(update_image(sample)) if bimpy.button('NEXT'): latents = rnd.randn(1, cfg.MODEL.LATENT_SPACE_SIZE) sample = torch.tensor(latents).float().cuda() im = bimpy.Image(update_image(sample)) if bimpy.slider_int("Layers", layers, 0, 8): im = bimpy.Image(update_image(sample)) bimpy.image(im, bimpy.Vec2(1024, 1024))
def train(cfg, logger, gpu_id=0): torch.cuda.set_device(gpu_id) model = Model( startf=cfg.MODEL.START_CHANNEL_COUNT, layer_count=cfg.MODEL.LAYER_COUNT, maxf=cfg.MODEL.MAX_CHANNEL_COUNT, latent_size=cfg.MODEL.LATENT_SPACE_SIZE, dlatent_avg_beta=cfg.MODEL.DLATENT_AVG_BETA, style_mixing_prob=cfg.MODEL.STYLE_MIXING_PROB, mapping_layers=cfg.MODEL.MAPPING_LAYERS, channels=3) model.cuda(gpu_id) model.train() if gpu_id == 0: model_s = Model( startf=cfg.MODEL.START_CHANNEL_COUNT, layer_count=cfg.MODEL.LAYER_COUNT, maxf=cfg.MODEL.MAX_CHANNEL_COUNT, latent_size=cfg.MODEL.LATENT_SPACE_SIZE, truncation_psi=cfg.MODEL.TRUNCATIOM_PSI, truncation_cutoff=cfg.MODEL.TRUNCATIOM_CUTOFF, mapping_layers=cfg.MODEL.MAPPING_LAYERS, channels=3) del model_s.discriminator model_s.cuda(gpu_id) model_s.eval() model_s.requires_grad_(False) generator = model.generator discriminator = model.discriminator mapping = model.mapping dlatent_avg = model.dlatent_avg count_param_override.print = lambda a: logger.info(a) logger.info("Trainable parameters generator:") count_parameters(generator) logger.info("Trainable parameters discriminator:") count_parameters(discriminator) generator_optimizer = LREQAdam([ {'params': generator.parameters()}, {'params': mapping.parameters()}], lr=cfg.TRAIN.BASE_LEARNING_RATE, betas=(cfg.TRAIN.ADAM_BETA_0, cfg.TRAIN.ADAM_BETA_1), weight_decay=0) discriminator_optimizer = LREQAdam([ {'params': discriminator.parameters()}, ], lr=cfg.TRAIN.BASE_LEARNING_RATE, betas=(cfg.TRAIN.ADAM_BETA_0, cfg.TRAIN.ADAM_BETA_1), weight_decay=0) scheduler = ComboMultiStepLR(optimizers={'generator': generator_optimizer,'discriminator': discriminator_optimizer}, milestones=cfg.TRAIN.LEARNING_DECAY_STEPS, # [] gamma=cfg.TRAIN.LEARNING_DECAY_RATE, # 0.1 reference_batch_size=32, base_lr=cfg.TRAIN.LEARNING_RATES) # 0.002 model_dict = { 'discriminator': discriminator, 'generator': generator, 'mapping': mapping, 'dlatent_avg': dlatent_avg } if gpu_id == 0: model_dict['generator_s'] = model_s.generator model_dict['mapping_s'] = model_s.mapping tracker = LossTracker(cfg.OUTPUT_DIR) checkpointer = Checkpointer(cfg, model_dict, { 'generator_optimizer': generator_optimizer, 'discriminator_optimizer': discriminator_optimizer, 'scheduler': scheduler, 'tracker': tracker }, logger=logger, save=gpu_id == 0) checkpointer.load() logger.info("Starting from epoch: %d" % (scheduler.start_epoch())) layer_to_resolution = generator.layer_to_resolution #[4, 8, 16, 32, 64, 128] dataset = TFRecordsDataset(cfg, logger, buffer_size_mb=1024) rnd = np.random.RandomState(3456) latents = rnd.randn(32, cfg.MODEL.LATENT_SPACE_SIZE) sample = torch.tensor(latents).float().cuda() lod2batch = lod_driver.LODDriver(cfg, logger, gpu_num=1, dataset_size=len(dataset)) #一个可以返回各类训练参数(param)的对象 for epoch in range(scheduler.start_epoch(), cfg.TRAIN.TRAIN_EPOCHS): model.train() lod2batch.set_epoch(epoch, [generator_optimizer, discriminator_optimizer]) logger.info("Batch size: %d, Batch size per GPU: %d, LOD: %d - %dx%d, blend: %.3f, dataset size: %d" % ( lod2batch.get_batch_size(), lod2batch.get_per_GPU_batch_size(), lod2batch.lod, 2 ** lod2batch.get_lod_power2(), 2 ** lod2batch.get_lod_power2(), lod2batch.get_blend_factor(), len(dataset))) dataset.reset(lod2batch.get_lod_power2(), lod2batch.get_per_GPU_batch_size()) batches = make_dataloader(cfg, logger, dataset, lod2batch.get_per_GPU_batch_size(), gpu_id) # 一个数据集分为多个batch,一个batch有n长图片 scheduler.set_batch_size(lod2batch.get_batch_size(), lod2batch.lod) #报错! need_permute = False with torch.autograd.profiler.profile(use_cuda=True, enabled=False) as prof: for x_orig in tqdm(batches): # x_orig:[-1,c,w,h] with torch.no_grad(): if x_orig.shape[0] != lod2batch.get_per_GPU_batch_size(): continue if need_permute: x_orig = x_orig.permute(0, 3, 1, 2) x_orig = (x_orig / 127.5 - 1.) blend_factor = lod2batch.get_blend_factor() needed_resolution = layer_to_resolution[lod2batch.lod] x = x_orig if lod2batch.in_transition: needed_resolution_prev = layer_to_resolution[lod2batch.lod - 1] x_prev = F.avg_pool2d(x_orig, 2, 2) x_prev_2x = F.interpolate(x_prev, needed_resolution) x = x * blend_factor + x_prev_2x * (1.0 - blend_factor) x.requires_grad = True discriminator_optimizer.zero_grad() loss_d = model(x, lod2batch.lod, blend_factor, d_train=True) tracker.update(dict(loss_d=loss_d)) loss_d.backward() discriminator_optimizer.step() if gpu_id == 0: betta = 0.5 ** (lod2batch.get_batch_size() / (10 * 1000.0)) model_s.lerp(model, betta) generator_optimizer.zero_grad() loss_g = model(x, lod2batch.lod, blend_factor, d_train=False) tracker.update(dict(loss_g=loss_g)) loss_g.backward() generator_optimizer.step() lod2batch.step() if gpu_id == 0: if lod2batch.is_time_to_save(): checkpointer.save("model_tmp_intermediate") if lod2batch.is_time_to_report(): save_sample(lod2batch, tracker, sample, x, logger, model_s, cfg, discriminator_optimizer, generator_optimizer) scheduler.step() if gpu_id == 0: checkpointer.save("model_tmp") save_sample(lod2batch, tracker, sample, x, logger, model_s, cfg, discriminator_optimizer, generator_optimizer) logger.info("Training finish!... save training results") if gpu_id == 0: checkpointer.save("model_final").wait()
"--config-file", default="configs/cat-bedroom-256.yaml", metavar="FILE", type=str, ) # args.config_file args = parser.parse_args() cfg = get_cfg_defaults() cfg.merge_from_file(args.config_file) cfg.freeze() #---------------------------------继承原版的Gs Gm , dlantent_avg model = Model( startf=cfg.MODEL.START_CHANNEL_COUNT, # startf = 16 layer_count=cfg.MODEL.LAYER_COUNT, # LAYER_COUNT: 9 maxf=cfg.MODEL.MAX_CHANNEL_COUNT, # cfg.MODEL.MAX_CHANNEL_COUNT : 512 latent_size=cfg.MODEL.LATENT_SPACE_SIZE, truncation_psi=cfg.MODEL.TRUNCATIOM_PSI, truncation_cutoff=cfg.MODEL. TRUNCATIOM_CUTOFF, # _C.MODEL.TRUNCATIOM_CUTOFF = 8 mapping_layers=cfg.MODEL.MAPPING_LAYERS, channels=3) model_dict = { 'generator_s': model.generator, 'mapping_fl_s': model.mapping, 'dlatent_avg': model.dlatent_avg, # dlatent_avg 平均人脸的位置 } #print('model.dlatent_avg_1:'+str(model.dlatent_avg.buff.data)) #[18,512],中心变量, 默认为0 logger = logging.getLogger("logger") checkpointer = Checkpointer(cfg, model_dict, logger=logger, save=True) checkpointer.load(
def load_from(name, cfg): dnnlib.tflib.init_tf() with open(name, 'rb') as f: m = pickle.load(f) Gs = m[2] #Gs_ = tflib.Network('G', func_name='stylegan.training.networks_stylegan.G_style', num_channels=3, resolution=1024) #Gs_.copy_vars_from(Gs) model = Model( startf=cfg.MODEL.START_CHANNEL_COUNT, layer_count=cfg.MODEL.LAYER_COUNT, maxf=cfg.MODEL.MAX_CHANNEL_COUNT, latent_size=cfg.MODEL.LATENT_SPACE_SIZE, mapping_layers=cfg.MODEL.MAPPING_LAYERS, truncation_psi=0.7, #cfg.MODEL.TRUNCATIOM_PSI, truncation_cutoff=cfg.MODEL.TRUNCATIOM_CUTOFF, channels=3) def tensor(x, transpose=None): x = Gs.vars[x].eval() if transpose: x = np.transpose(x, transpose) return torch.tensor(x) for i in range(cfg.MODEL.MAPPING_LAYERS): block = getattr(model.mapping, "block_%d" % (i + 1)) block.fc.weight[:] = tensor('G_mapping/Dense%d/weight' % i, (1, 0)) * block.fc.std block.fc.bias[:] = tensor( 'G_mapping/Dense%d/bias' % i) * block.fc.lrmul model.dlatent_avg.buff[:] = tensor('dlatent_avg') model.generator.const[:] = tensor('G_synthesis/4x4/Const/const') for i in range(model.generator.layer_count): j = model.generator.layer_count - i - 1 name = '%dx%d' % (2**(2 + i), 2**(2 + i)) block = model.generator.decode_block[i] prefix = 'G_synthesis/%s' % name if not block.has_first_conv: prefix_1 = '%s/Const' % prefix prefix_2 = '%s/Conv' % prefix else: prefix_1 = '%s/Conv0_up' % prefix prefix_2 = '%s/Conv1' % prefix block.noise_weight_1[0, :, 0, 0] = tensor('%s/Noise/weight' % prefix_1) block.noise_weight_2[0, :, 0, 0] = tensor('%s/Noise/weight' % prefix_2) if block.has_first_conv: if block.fused_scale: block.conv_1.weight[:] = tensor( '%s/weight' % prefix_1, (2, 3, 0, 1)) * block.conv_1.std else: block.conv_1.weight[:] = tensor( '%s/weight' % prefix_1, (3, 2, 0, 1)) * block.conv_1.std block.conv_2.weight[:] = tensor('%s/weight' % prefix_2, (3, 2, 0, 1)) * block.conv_2.std block.bias_1[0, :, 0, 0] = tensor('%s/bias' % prefix_1) block.bias_2[0, :, 0, 0] = tensor('%s/bias' % prefix_2) block.style_1.weight[:] = tensor('%s/StyleMod/weight' % prefix_1, (1, 0)) * block.style_1.std block.style_1.bias[:] = tensor('%s/StyleMod/bias' % prefix_1) block.style_2.weight[:] = tensor('%s/StyleMod/weight' % prefix_2, (1, 0)) * block.style_2.std block.style_2.bias[:] = tensor('%s/StyleMod/bias' % prefix_2) model.generator.to_rgb[i].to_rgb.weight[:] = tensor( 'G_synthesis/ToRGB_lod%d/weight' % (j), (3, 2, 0, 1)) * model.generator.to_rgb[i].to_rgb.std model.generator.to_rgb[i].to_rgb.bias[:] = tensor( 'G_synthesis/ToRGB_lod%d/bias' % (j)) return model, Gs
def parallel_regularized_evolution(profiler, args): population = collections.deque() # kill_over_time = dict() scheduler = AsyncScheduler(profiler, n_workers=args.n_workers) tracker = Tracker() skipped_archs = 0 non_skipped_archs = 0 if args.predict_model: dp = DataPreparer(args.benchmark) predictor = Predictor() step = 0 # Initialize the population with random models. init_size = args.sample_size while step < init_size: model = Model() model.arch = ModelHelper.random_cell(profiler) while scheduler.full(): m = scheduler.fetch_next_evaluating_model() tracker.record_model(m) population.append(m) scheduler.parallel_train_and_evaluate(model) step += 1 finished_models = scheduler.wait_for_finish() tracker.record_model(finished_models) population += finished_models if args.predict_model: dp.load_evaluated_models(tracker.evaluated_models) predictor.fit(dp) cached_models = [] while step < args.n_iters: # print(step, 'Pop ({}) Accuracies'.format(len(population)), [model.accuracy for model in population][:10]) if scheduler.full(): evaluated_model = scheduler.fetch_next_evaluating_model() tracker.record_model(evaluated_model) #update predictor if args.predict_model: if not evaluated_model.es: cached_models.append(evaluated_model) if len(cached_models) == 10: dp.load_evaluated_models(cached_models) cached_models = [] predictor.fit(dp) ## update population if not evaluated_model.es: population.append(evaluated_model) while len(population) > args.pop_size: if args.predict_kill: # Sample randomly chosen models from the current population. sample = [] while len(sample) < args.sample_size: # Inefficient, but written this way for clarity. In the case of neural # nets, the efficiency of this line is irrelevant because training neural # nets is the rate-determining step. candidate = random.choice(list(population)) sample.append(candidate) predicted_accuracies = [] for each in sample: candidates_acc = [ predictor.predict( np.array([ dp.json_model_feature(each.serialize()) ]))[0] ] # candidates_acc = [each.accuracy] while True: c = Model() c.arch = ModelHelper.mutate( each, args.mutate_step, profiler) result = profiler.profile( c, return_fixed_metrics=True) c.training_time = result['training_time'] c.cost = result['cost'] c.params = result['trainable_parameters'] accuracy = predictor.predict( np.array( [dp.json_model_feature(c.serialize())]))[0] candidates_acc.append(accuracy) if len(candidates_acc) == args.kill_num: break arch_predicted_acc = np.mean( candidates_acc ) + args.kill_alpha * np.std(candidates_acc) predicted_accuracies.append(arch_predicted_acc) min_id = np.argmin(predicted_accuracies) # if step % 20 == 0: # kill_over_time[step] = ([sample[min_id].accuracy], [sample[min_id].test_accuracy]) population.remove(sample[min_id]) else: if args.method == 're': population.popleft() while not scheduler.full(): # fill them # Sample randomly chosen models from the current population. sample = [] while len(sample) < args.sample_size: # Inefficient, but written this way for clarity. In the case of neural # nets, the efficiency of this line is irrelevant because training neural # nets is the rate-determining step. candidate = random.choice(list(population)) sample.append(candidate) predicted_accuracies = [] for each in sample: candidates_acc = [ predictor.predict( np.array([dp.json_model_feature(each.serialize())]))[0] ] # candidates_acc = [each.accuracy] while True: c = Model() c.arch = ModelHelper.mutate(each, args.mutate_step, profiler) result = profiler.profile(c, return_fixed_metrics=True) c.training_time = result['training_time'] c.cost = result['cost'] c.params = result['trainable_parameters'] accuracy = predictor.predict( np.array([dp.json_model_feature(c.serialize())]))[0] candidates_acc.append(accuracy) if len(candidates_acc) == args.kill_num: break arch_predicted_acc = np.mean( candidates_acc) + args.kill_alpha * np.std(candidates_acc) predicted_accuracies.append(arch_predicted_acc) max_id = np.argmax(predicted_accuracies) parent = sample[max_id] if args.predict_mutate: candidate_children = [] while True: c = Model() c.arch = ModelHelper.mutate(parent, args.mutate_step, profiler) result = profiler.profile(c, return_fixed_metrics=True) c.training_time = result['training_time'] c.cost = result['cost'] c.params = result['trainable_parameters'] c.accuracy = predictor.predict( np.array([dp.json_model_feature(c.serialize())]))[0] candidate_children.append(c) if len(candidate_children) == args.mutate_num: break child = max(candidate_children, key=lambda i: i.fitness(mo=args.mo)) else: child = Model() child.arch = ModelHelper.mutate(parent, args.mutate_step, profiler) child.parent = parent child.root = parent.root if parent.root else parent scheduler.parallel_train_and_evaluate(child) step += 1 if step % 1000 == 0: print('step', step, 'population', stats.describe([m.accuracy for m in population])) finished_models = scheduler.wait_for_finish() tracker.record_model(finished_models) return tracker