def run(cfg, model: BaseModel, dataset: BaseDataset, device, measurement_name: str): measurements = {} num_batches = getattr(cfg.debugging, "num_batches", np.inf) run_epoch(model, dataset.train_dataloader(), device, num_batches) measurements["train"] = extract_histogram(model.get_spatial_ops(), normalize=False) if dataset.has_val_loader: run_epoch(model, dataset.val_dataloader(), device, num_batches) measurements["val"] = extract_histogram(model.get_spatial_ops(), normalize=False) for loader_idx, loader in enumerate(dataset.test_dataloaders()): run_epoch(model, dataset.test_dataloaders(), device, num_batches) measurements[dataset.get_test_dataset_name( loader_idx)] = extract_histogram(model.get_spatial_ops(), normalize=False) with open( os.path.join(DIR, "measurements/{}.pickle".format(measurement_name)), "wb") as f: pickle.dump(measurements, f)
def __init__(self, opt): """Initialize this model class. Parameters: opt -- training/test options A few things can be done here. - (required) call the initialization function of BaseModel - define loss function, visualization images, model names, and optimizers """ BaseModel.__init__(self, opt) # call the initialization method of BaseModel # specify the training losses you want to print out. The program will call base_model.get_current_losses to plot the losses to the console and save them to the disk. self.loss_names = ['loss_G'] # specify the images you want to save and display. The program will call base_model.get_current_visuals to save and display these images. self.visual_names = ['data_A', 'data_B', 'output'] # specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks to save and load networks. # you can use opt.isTrain to specify different behaviors for training and test. For example, some networks will not be used during test, and you don't need to load them. self.model_names = ['G'] # define networks; you can use opt.isTrain to specify different behaviors for training and test. self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, gpu_ids=self.gpu_ids) if self.isTrain: # only defined during training time # define your loss functions. You can use losses provided by torch.nn such as torch.nn.L1Loss. # We also provide a GANLoss class "networks.GANLoss". self.criterionGAN = networks.GANLoss().to(self.device) self.criterionLoss = torch.nn.L1Loss() # define and initialize optimizers. You can define one optimizer for each network. # If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example. self.optimizer = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers = [self.optimizer]
def initialize_model(self, model: BaseModel, weight_name: str = None): if not self._checkpoint.is_empty: state_dict = self._checkpoint.get_state_dict(weight_name) model.load_state_dict(state_dict) optimizer = self._checkpoint.get_optimizer() lr_params = self._checkpoint.get_lr_params() model.set_optimizer(optimizer.__class__, lr_params=lr_params)
def _initialize_model(self, model: BaseModel, weight_name): if not self._checkpoint.is_empty: state_dict = self._checkpoint.get_state_dict(weight_name) model.load_state_dict(state_dict) if self._resume: model.optimizer = self._checkpoint.get_optimizer(model) model.schedulers = self._checkpoint.get_schedulers(model)
def run_epoch(model: BaseModel, loader, device: str, num_batches: int): model.eval() with Ctq(loader) as tq_loader: for batch_idx, data in enumerate(tq_loader): if batch_idx < num_batches: process(model, data, device) else: break
def eval_epoch(model: BaseModel, dataset, device, tracker: BaseTracker, checkpoint: ModelCheckpoint): tracker.reset("val") loader = dataset.val_dataloader with Ctq(loader) as tq_val_loader: for data in tq_val_loader: with torch.no_grad(): model.set_input(data, device) model.forward() tracker.track(model) tq_val_loader.set_postfix(**tracker.get_metrics(), color=COLORS.VAL_COLOR) tracker.print_summary()
def test_epoch(model: BaseModel, dataset, device, tracker: BaseTracker, checkpoint: ModelCheckpoint): tracker.reset("test") loader = dataset.test_dataloader() with Ctq(loader) as tq_test_loader: for data in tq_test_loader: data = data.to(device) with torch.no_grad(): model.set_input(data) model.forward() tracker.track(model) tq_test_loader.set_postfix(**tracker.get_metrics(), color=COLORS.TEST_COLOR) tracker.print_summary()
def track(self, model: BaseModel): """ Add current model predictions (usually the result of a batch) to the tracking """ super().track(model) outputs = self._convert(model.get_output()) targets = self._convert(model.get_labels()) assert outputs.shape[0] == len(targets) self._confusion_matrix.count_predicted_batch(targets, np.argmax(outputs, 1)) self._acc = 100 * self._confusion_matrix.get_overall_accuracy() self._macc = 100 * self._confusion_matrix.get_mean_class_accuracy() self._miou = 100 * self._confusion_matrix.get_average_intersection_union( )
def track(self, model: BaseModel): """ Add current model predictions (usually the result of a batch) to the tracking """ super().track(model) outputs = self._convert(model.get_output()) targets = self._convert(model.get_labels()) erp = torch.sqrt(((outputs - targets) / (targets + self._eps)) ** 2) self._merp = torch.mean(erp).item() self._mer = ( torch.mean(F.normalize(outputs - targets, p=2, dim=-1)) / torch.mean((F.normalize(targets, p=2, dim=-1) + self._eps)) ).item()
def run(model: BaseModel, dataset: BaseDataset, device, output_path, cfg): # Set dataloaders num_fragment = dataset.num_fragment if cfg.data.is_patch: for i in range(num_fragment): dataset.set_patches(i) dataset.create_dataloaders( model, cfg.batch_size, False, cfg.num_workers, False, ) loader = dataset.test_dataloaders()[0] features = [] scene_name, pc_name = dataset.get_name(i) with Ctq(loader) as tq_test_loader: for data in tq_test_loader: # pcd = open3d.geometry.PointCloud() # pcd.points = open3d.utility.Vector3dVector(data.pos[0].numpy()) # open3d.visualization.draw_geometries([pcd]) with torch.no_grad(): model.set_input(data, device) model.forward() features.append(model.get_output().cpu()) features = torch.cat(features, 0).numpy() log.info("save {} from {} in {}".format(pc_name, scene_name, output_path)) save(output_path, scene_name, pc_name, dataset.base_dataset[i].to("cpu"), features) else: dataset.create_dataloaders( model, 1, False, cfg.num_workers, False, ) loader = dataset.test_dataloaders()[0] with Ctq(loader) as tq_test_loader: for i, data in enumerate(tq_test_loader): with torch.no_grad(): model.set_input(data, device) model.forward() features = model.get_output()[0] # batch of 1 save(output_path, scene_name, pc_name, data.to("cpu"), features)
def track(self, model: BaseModel): """ Add model predictions (accuracy) """ super().track(model) outputs = self._convert(model.get_output()) self._acc = compute_accuracy(outputs[::2], outputs[1::2])
def test_epoch(model: BaseModel, dataset, device, tracker: BaseTracker, checkpoint: ModelCheckpoint): loaders = dataset.test_dataloaders for loader in loaders: stage_name = loader.dataset.name tracker.reset(stage_name) with Ctq(loader) as tq_test_loader: for data in tq_test_loader: with torch.no_grad(): model.set_input(data, device) model.forward() tracker.track(model) tq_test_loader.set_postfix(**tracker.get_metrics(), color=COLORS.TEST_COLOR) tracker.print_summary()
def run(model: BaseModel, dataset: BaseDataset, device, output_path): loaders = dataset.test_dataloaders predicted = {} for loader in loaders: loader.dataset.name with Ctq(loader) as tq_test_loader: for data in tq_test_loader: with torch.no_grad(): model.set_input(data, device) model.forward() predicted = { **predicted, **dataset.predict_original_samples(data, model.conv_type, model.get_output()) } save(output_path, predicted)
def eval_epoch( epoch: int, model: BaseModel, dataset, device, tracker: BaseTracker, checkpoint: ModelCheckpoint, visualizer: Visualizer, early_break: bool, ): model.eval() tracker.reset("val") visualizer.reset(epoch, "val") loader = dataset.val_dataloader() with Ctq(loader) as tq_val_loader: for data in tq_val_loader: data = data.to(device) with torch.no_grad(): model.set_input(data) model.forward() tracker.track(model) tq_val_loader.set_postfix(**tracker.get_metrics(), color=COLORS.VAL_COLOR) if visualizer.is_active: visualizer.save_visuals(model.get_current_visuals()) if early_break: break metrics = tracker.publish(epoch) tracker.print_summary() checkpoint.save_best_models_under_current_metrics(model, metrics)
def run(model: BaseModel, dataset: BaseDataset, device, output_path): loaders = dataset.test_dataloaders() predicted = {} for idx, loader in enumerate(loaders): dataset.get_test_dataset_name(idx) with Ctq(loader) as tq_test_loader: for data in tq_test_loader: data = data.to(device) with torch.no_grad(): model.set_input(data) model.forward() predicted = { **predicted, **dataset.predict_original_samples(data, model.conv_type, model.get_output()) } save(output_path, predicted)
def save_best_models_under_current_metrics(self, model: BaseModel, metrics_holder: dict, **kwargs): """[This function is responsible to save checkpoint under the current metrics and their associated DEFAULT_METRICS_FUNC] Arguments: model {[BaseModel]} -- [Model] metrics_holder {[Dict]} -- [Need to contain stage, epoch, current_metrics] """ metrics = metrics_holder["current_metrics"] stage = metrics_holder["stage"] epoch = metrics_holder["epoch"] stats = self._checkpoint.stats state_dict = copy.deepcopy(model.state_dict()) current_stat = {} current_stat["epoch"] = epoch models_to_save = self._checkpoint.models if stage not in stats: stats[stage] = [] if stage == "train": models_to_save[Checkpoint._LATEST] = state_dict else: if len(stats[stage]) > 0: latest_stats = stats[stage][-1] msg = "" improved_metric = 0 for metric_name, current_metric_value in metrics.items(): current_stat[metric_name] = current_metric_value metric_func = self.find_func_from_metric_name(metric_name, DEFAULT_METRICS_FUNC) best_metric_from_stats = latest_stats.get("best_{}".format(metric_name), current_metric_value) best_value = metric_func(best_metric_from_stats, current_metric_value) current_stat["best_{}".format(metric_name)] = best_value # This new value seems to be better under metric_func if (self._selection_stage == stage) and ( current_metric_value == best_value ): # Update the model weights models_to_save["best_{}".format(metric_name)] = state_dict msg += "{}: {} -> {}, ".format(metric_name, best_metric_from_stats, best_value) improved_metric += 1 if improved_metric > 0: colored_print(COLORS.VAL_COLOR, msg[:-2]) else: # stats[stage] is empty. for metric_name, metric_value in metrics.items(): current_stat[metric_name] = metric_value current_stat["best_{}".format(metric_name)] = metric_value models_to_save["best_{}".format(metric_name)] = state_dict self._checkpoint.stats[stage].append(current_stat) self._checkpoint.save_objects(models_to_save, stage, current_stat, model.optimizer, model.schedulers, **kwargs)
def track(self, model: BaseModel): """ Add model predictions (accuracy) """ super().track(model) outputs = self._convert(model.get_output()) N = len(outputs) // 2 self._acc = compute_accuracy(outputs[:N], outputs[N:])
def __init__(self, opt): """Initialize the pix2pixPL class. Parameters: opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions """ BaseModel.__init__(self, opt) # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses> self.loss_names = ['G_GAN', 'G_Style', 'G_Content', 'D_real', 'D_fake'] # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals> self.visual_names = ['real_A', 'fake_B', 'real_B'] # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks> if self.isTrain: self.model_names = ['G', 'D'] else: # during test time, only load G self.model_names = ['G'] # define networks (both generator and discriminator) self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids) if self.isTrain: # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids) if self.isTrain: # define loss functions self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device) #self.criterionL1 = torch.nn.L1Loss() self.criterionPL = PerceptualLoss(self.device).to(self.device) # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>. self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D)
def train_epoch(epoch, model: BaseModel, dataset, device: str, tracker: BaseTracker, checkpoint: ModelCheckpoint, log): model.train() tracker.reset("train") train_loader = dataset.train_dataloader() iter_data_time = time.time() with Ctq(train_loader) as tq_train_loader: for i, data in enumerate(tq_train_loader): data = data.to(device) # This takes time model.set_input(data) t_data = time.time() - iter_data_time iter_start_time = time.time() model.optimize_parameters(dataset.batch_size) if i % 10 == 0: tracker.track(model) tq_train_loader.set_postfix(**tracker.get_metrics(), data_loading=float(t_data), iteration=float(time.time() - iter_start_time), color=COLORS.TRAIN_COLOR) iter_data_time = time.time() metrics = tracker.publish() checkpoint.save_best_models_under_current_metrics(model, metrics) log.info("Learning rate = %f" % model.learning_rate)
def __init__(self, opt): """Initialize the pix2pix class. Parameters: opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions """ assert (not opt.isTrain) BaseModel.__init__(self, opt) # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses> self.loss_names = [] # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals> self.visual_names = ['real', 'fake'] # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks> self.model_names = ['G' + opt.model_suffix ] # only generator is needed. self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids) # assigns the model to self.netG_[suffix] so that it can be loaded # please see <BaseModel.load_networks> setattr(self, 'netG' + opt.model_suffix, self.netG) # store netG in self.
def _init_from_compact_format(self, opt, model_type, dataset, modules_lib): """Create a unetbasedmodel from the compact options format - where the same convolution is given for each layer, and arguments are given in lists """ self.down_modules = nn.ModuleList() self.inner_modules = nn.ModuleList() self.up_modules = nn.ModuleList() # Factory for creating up and down modules factory_module_cls = self._get_factory(model_type, modules_lib) down_conv_cls_name = opt.down_conv.module_name up_conv_cls_name = opt.up_conv.module_name self._factory_module = factory_module_cls( down_conv_cls_name, up_conv_cls_name, modules_lib) # Create the factory object # Loal module contains_global = hasattr(opt, "innermost") and opt.innermost is not None if contains_global: inners = self._create_inner_modules(opt.innermost, modules_lib) for inner in inners: self.inner_modules.append(inner) else: self.inner_modules.append(Identity()) # Down modules for i in range(len(opt.down_conv.down_conv_nn)): args = self._fetch_arguments(opt.down_conv, i, "DOWN") conv_cls = self._get_from_kwargs(args, "conv_cls") down_module = conv_cls(**args) self._save_sampling_and_search(down_module) self.down_modules.append(down_module) # Up modules for i in range(len(opt.up_conv.up_conv_nn)): args = self._fetch_arguments(opt.up_conv, i, "UP") conv_cls = self._get_from_kwargs(args, "conv_cls") up_module = conv_cls(**args) self._save_upsample(up_module) self.up_modules.append(up_module) self.metric_loss_module, self.miner_module = BaseModel.get_metric_loss_and_miner( getattr(opt, "loss", None), getattr(opt, "miner", None))
def train_epoch( epoch: int, model: BaseModel, dataset, device: str, tracker: BaseTracker, checkpoint: ModelCheckpoint, visualizer: Visualizer, early_break: bool, ): model.train() tracker.reset("train") visualizer.reset(epoch, "train") train_loader = dataset.train_dataloader iter_data_time = time.time() with Ctq(train_loader) as tq_train_loader: for i, data in enumerate(tq_train_loader): model.set_input(data, device) t_data = time.time() - iter_data_time iter_start_time = time.time() model.optimize_parameters(epoch, dataset.batch_size) if i % 10 == 0: tracker.track(model) tq_train_loader.set_postfix(**tracker.get_metrics(), data_loading=float(t_data), iteration=float(time.time() - iter_start_time), color=COLORS.TRAIN_COLOR) if visualizer.is_active: visualizer.save_visuals(model.get_current_visuals()) iter_data_time = time.time() if early_break: break metrics = tracker.publish(epoch) checkpoint.save_best_models_under_current_metrics(model, metrics) log.info("Learning rate = %f" % model.learning_rate)
def _init_from_compact_format(self, opt, model_type, dataset, modules_lib): """Create a backbonebasedmodel from the compact options format - where the same convolution is given for each layer, and arguments are given in lists """ num_convs = len(opt.down_conv.down_conv_nn) self.down_modules = nn.ModuleList() factory_module_cls = self._get_factory(model_type, modules_lib) down_conv_cls_name = opt.down_conv.module_name self._factory_module = factory_module_cls(down_conv_cls_name, None, modules_lib) # Down modules for i in range(num_convs): args = self._fetch_arguments(opt.down_conv, i, "DOWN") conv_cls = self._get_from_kwargs(args, "conv_cls") down_module = conv_cls(**args) self._save_sampling_and_search(down_module) self.down_modules.append(down_module) self.metric_loss_module, self.miner_module = BaseModel.get_metric_loss_and_miner( getattr(opt, "metric_loss", None), getattr(opt, "miner", None))
def eval_epoch(model: BaseModel, dataset, device, tracker: BaseTracker, checkpoint: ModelCheckpoint, log): model.eval() tracker.reset("val") loader = dataset.val_dataloader() with Ctq(loader) as tq_val_loader: for data in tq_val_loader: data = data.to(device) with torch.no_grad(): model.set_input(data) model.forward() tracker.track(model) tq_val_loader.set_postfix(**tracker.get_metrics(), color=COLORS.VAL_COLOR) metrics = tracker.publish() tracker.print_summary() checkpoint.save_best_models_under_current_metrics(model, metrics)
def test_epoch( epoch: int, model: BaseModel, dataset, device, tracker: BaseTracker, checkpoint: ModelCheckpoint, visualizer: Visualizer, early_break: bool, ): model.eval() loaders = dataset.test_dataloaders() for idx, loader in enumerate(loaders): stage_name = dataset.get_test_dataset_name(idx) tracker.reset(stage_name) visualizer.reset(epoch, stage_name) with Ctq(loader) as tq_test_loader: for data in tq_test_loader: data = data.to(device) with torch.no_grad(): model.set_input(data) model.forward() tracker.track(model) tq_test_loader.set_postfix(**tracker.get_metrics(), color=COLORS.TEST_COLOR) if visualizer.is_active: visualizer.save_visuals(model.get_current_visuals()) if early_break: break metrics = tracker.publish(epoch) tracker.print_summary() checkpoint.save_best_models_under_current_metrics(model, metrics)