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)
