def main(log_dir: str):
for path in tqdm.tqdm(glob.glob(log_dir)):
experiment_id = None
first_dev = None
full_dev = None
with open(path) as f:
for line in f:
maybe_first = re.match(FIRST_RE, line)
if maybe_first is not None:
first_dev = float(maybe_first.group(1))
maybe_full = re.match(FULL_RE, line)
if maybe_full is not None:
full_dev = float(maybe_full.group(1))
maybe_id = re.match(EXP_RE, line)
if maybe_id is not None:
experiment_id = maybe_id.group(1)
if experiment_id is not None and first_dev is not None and full_dev is not None:
print(experiment_id, first_dev, full_dev)
experiment = comet_ml.ExistingExperiment(
previous_experiment=experiment_id)
experiment.log_metric("dev_first_accuracy", first_dev)
experiment.log_metric("dev_full_accuracy", full_dev)
experiment.end()
def _get_experiment(self, mode, experiment_id=None):
if mode == "offline":
if experiment_id is not None:
return comet_ml.ExistingOfflineExperiment(
previous_experiment=experiment_id,
workspace=self.workspace,
project_name=self.project_name,
**self.experiment_kwargs,
)
return comet_ml.OfflineExperiment(
workspace=self.workspace,
project_name=self.project_name,
**self.experiment_kwargs,
)
else:
if experiment_id is not None:
return comet_ml.ExistingExperiment(
previous_experiment=experiment_id,
workspace=self.workspace,
project_name=self.project_name,
**self.experiment_kwargs,
)
return comet_ml.Experiment(
workspace=self.workspace,
project_name=self.project_name,
**self.experiment_kwargs,
)
def score_model(serialization_dir: Text,
log_to_comet=False,
comet_experiment_id: Optional[str] = None):
archive = load_archive(os.path.join(serialization_dir, "model.tar.gz"),
cuda_device=0)
predictor = QbPredictor.from_archive(
archive,
predictor_name="qb.predictor.QbPredictor",
)
dataset_reader = predictor._dataset_reader # pylint: disable=protected-access
dataset_reader.include_label = False
dataset_reader.first_sentence_only = True
dataset_reader.full_question_only = False
dev_first_sentence = dataset_reader.read("guessdev")
accuracy_start_dev = compute_accuracy(predictor, dev_first_sentence)
log.info("First dev accuracy: %s", accuracy_start_dev)
dataset_reader.first_sentence_only = False
dataset_reader.full_question_only = True
dev_full_question = dataset_reader.read("guessdev")
accuracy_full_dev = compute_accuracy(predictor, dev_full_question)
log.info("Full dev accuracy: %s", accuracy_full_dev)
log.info("log_to_comet: %s", log_to_comet)
if log_to_comet:
if comet_experiment_id is None:
experiment = comet_ml.get_global_experiment()
else:
experiment = comet_ml.ExistingExperiment(
previous_experiment=comet_experiment_id)
experiment.log_metric("dev_first_accuracy", accuracy_start_dev)
experiment.log_metric("dev_full_accuracy", accuracy_full_dev)
def main():
args = parse_args()
if os.path.exists(args.output_path):
print(f"Output path {args.output_path} exists!!!")
return
random.seed(args.seed)
np.random.seed(args.seed)
if args.limit:
print(
"WARNING: --limit SHOULD ONLY BE USED FOR TESTING. REAL METRICS SHOULD NOT BE COMPUTED USING LIMIT."
)
if args.tasks == "all_tasks":
task_names = tasks.ALL_TASKS
else:
task_names = args.tasks.split(",")
task_dict = tasks.get_task_dict(task_names)
lm = models.get_model(args.model)
train_args = simple_parse_args_string(args.train_args)
model_args = simple_parse_args_string(args.model_args)
if train_args:
train_args.update(model_args)
train_args["seed"] = args.seed
results = evaluator.evaluate(lm, task_dict, args.provide_description,
args.num_fewshot, args.limit, train_args,
args.model_args, args.seed)
results["args"] = args.__dict__
dumped = json.dumps(results, indent=2)
print(dumped)
if args.output_path:
with open(args.output_path, "w") as f:
f.write(dumped)
for task, task_res in results.items():
if task not in task_names:
continue
if "train_args" not in task_res:
experiment = comet_ml.Experiment(
api_key=os.environ.get('COMET_API_KEY'),
project_name=os.environ.get('COMET_PROJECT', "few-shot"),
workspace=os.environ.get('COMET_WORKSPACE', "yuvalkirstain"),
)
experiment.log_asset(args.output_path)
else:
experiment = comet_ml.ExistingExperiment(
api_key=os.environ.get('COMET_API_KEY'),
previous_experiment=task_res["train_args"]
["previous_experiment"])
experiment.log_asset(args.output_path)
def predict(self, *args, **kwargs):
import comet_ml
try:
self.experiment = comet_ml.ExistingExperiment()
except Exception:
logging.error("Ignored --comet. No '.comet.config' file")
return
cli = self._make_command_line(args)
self._log_html(cli)
def test(self, *args, **kwargs):
import comet_ml
try:
self.cometml_experiment = comet_ml.ExistingExperiment()
except Exception:
logging.error("Ignored --comet. No '.comet.config' file")
return
logging.info("comet.test() called......")
cli = self._make_command_line(args)
self._log_html(cli)
def get_experiment(run_id):
experiment_id = hashlib.sha1(run_id.encode("utf-8")).hexdigest()
os.environ["COMET_EXPERIMENT_KEY"] = experiment_id
api = comet_ml.API() # Assumes API key is set in config/env
api_experiment = api.get_experiment_by_id(experiment_id)
if api_experiment is None:
return comet_ml.Experiment(project_name=PROJECT_NAME)
else:
return comet_ml.ExistingExperiment(project_name=PROJECT_NAME)
def init_logger(self, cfg):
logger = None
# Check to see if there is a key in environment:
EXPERIMENT_KEY = cfg.experiment_key
# First, let's see if we continue or start fresh:
CONTINUE_RUN = cfg.resume
if (EXPERIMENT_KEY is not None):
# There is one, but the experiment might not exist yet:
api = comet_ml.API() # Assumes API key is set in config/env
try:
api_experiment = api.get_experiment_by_id(EXPERIMENT_KEY)
except Exception:
api_experiment = None
if api_experiment is not None:
CONTINUE_RUN = True
# We can get the last details logged here, if logged:
# step = int(api_experiment.get_parameters_summary("batch")["valueCurrent"])
# epoch = int(api_experiment.get_parameters_summary("epochs")["valueCurrent"])
if CONTINUE_RUN:
# 1. Recreate the state of ML system before creating experiment
# otherwise it could try to log params, graph, etc. again
# ...
# 2. Setup the existing experiment to carry on:
logger = comet_ml.ExistingExperiment(
previous_experiment=EXPERIMENT_KEY,
log_env_details=True, # to continue env logging
log_env_gpu=True, # to continue GPU logging
log_env_cpu=True, # to continue CPU logging
auto_histogram_weight_logging=True,
auto_histogram_gradient_logging=True,
auto_histogram_activation_logging=True)
# Retrieved from above APIExperiment
# self.logger.set_epoch(epoch)
else:
# 1. Create the experiment first
# This will use the COMET_EXPERIMENT_KEY if defined in env.
# Otherwise, you could manually set it here. If you don't
# set COMET_EXPERIMENT_KEY, the experiment will get a
# random key!
logger = comet_ml.Experiment(
disabled=cfg.disabled,
project_name=cfg.project,
auto_histogram_weight_logging=True,
auto_histogram_gradient_logging=True,
auto_histogram_activation_logging=True)
logger.add_tags(cfg.tags.split())
logger.log_parameters(self.cfg)
return logger
def get_old_experiment(name, number=-1, key=None):
"""Get the Comet.ml experiment with the specified name and of the specified index
(out of all experiments with that name)
Keys should be stored in experiments/comet_keys.pkl
:param name: name of the Comet.ml experiment
:param number: list index of the needed experiment
:param key: if specified, just fetches an experiment with that key, ignoring name and index
:return:
"""
if key is None:
experiment_keys_path = pathlib.Path('experiments/comet_keys.pkl')
if not experiment_keys_path.exists():
raise Exception(
f'No experiment_keys file in {experiment_keys_path}')
with open(experiment_keys_path, 'rb') as f:
experiment_keys = pickle.load(f)
keys = experiment_keys[name]
if not keys:
raise Exception(f'No saved keys for experiment named {name}')
key = keys[number]
experiment = comet_ml.ExistingExperiment(previous_experiment=key,
display_summary_level=0)
return experiment
def attach_comet_ml_experiment(self, api_key, exp_key):
self._comet_experiment = comet_ml.ExistingExperiment(
api_key=api_key, previous_experiment=exp_key
)
def training_loop(
run_dir = '.', # Output directory.
training_set_kwargs = {}, # Options for training set.
validation_set_kwargs = {},
data_loader_kwargs = {}, # Options for torch.utils.data.DataLoader.
G_kwargs = {}, # Options for generator network.
D_kwargs = {}, # Options for discriminator network.
G_opt_kwargs = {}, # Options for generator optimizer.
D_opt_kwargs = {}, # Options for discriminator optimizer.
augment_kwargs = None, # Options for augmentation pipeline. None = disable.
loss_kwargs = {}, # Options for loss function.
gradient_clipping = None, # Max global gradient norm for clippingor None to disable it
metrics = [], # Metrics to evaluate during training.
comet_api_key = '', # comet_ml api key for comet.ml logging or '' to disable it
comet_experiment_key = '', # comet_ml experiment key for comet.ml logging or '' to disable it
random_seed = 0, # Global random seed.
num_gpus = 1, # Number of GPUs participating in the training.
rank = 0, # Rank of the current process in [0, num_gpus[.
batch_size = 4, # Total batch size for one training iteration. Can be larger than batch_gpu * num_gpus.
batch_gpu = 4, # Number of samples processed at a time by one GPU.
ema_kimg = 10, # Half-life of the exponential moving average (EMA) of generator weights.
ema_rampup = None, # EMA ramp-up coefficient.
G_reg_interval = 4, # How often to perform regularization for G? None = disable lazy regularization.
D_reg_interval = 16, # How often to perform regularization for D? None = disable lazy regularization.
augment_p = 0, # Initial value of augmentation probability.
ada_target = None, # ADA target value. None = fixed p.
ada_interval = 4, # How often to perform ADA adjustment?
ada_kimg = 500, # ADA adjustment speed, measured in how many kimg it takes for p to increase/decrease by one unit.
total_kimg = 25000, # Total length of the training, measured in thousands of real images.
kimg_per_tick = 4, # Progress snapshot interval.
image_snapshot_ticks = 50, # How often to save image snapshots? None = disable.
network_snapshot_ticks = 50, # How often to save network snapshots? None = disable.
resume_pkl = None, # Network pickle to resume training from.
cudnn_benchmark = True, # Enable torch.backends.cudnn.benchmark?
allow_tf32 = False, # Enable torch.backends.cuda.matmul.allow_tf32 and torch.backends.cudnn.allow_tf32?
abort_fn = None, # Callback function for determining whether to abort training. Must return consistent results across ranks.
progress_fn = None, # Callback function for updating training progress. Called for all ranks.
):
# Initialize.
start_time = time.time()
device = torch.device('cuda', rank)
np.random.seed(random_seed * num_gpus + rank)
torch.manual_seed(random_seed * num_gpus + rank)
torch.backends.cudnn.benchmark = cudnn_benchmark # Improves training speed.
torch.backends.cuda.matmul.allow_tf32 = allow_tf32 # Allow PyTorch to internally use tf32 for matmul
torch.backends.cudnn.allow_tf32 = allow_tf32 # Allow PyTorch to internally use tf32 for convolutions
conv2d_gradfix.enabled = True # Improves training speed.
grid_sample_gradfix.enabled = True # Avoids errors with the augmentation pipe.
# Load training set.
if rank == 0:
print('Loading training set...')
training_set = dnnlib.util.construct_class_by_name(**training_set_kwargs) # subclass of training.dataset.Dataset
training_set_sampler = misc.InfiniteSampler(dataset=training_set, rank=rank, num_replicas=num_gpus, seed=random_seed)
training_set_iterator = iter(torch.utils.data.DataLoader(dataset=training_set, sampler=training_set_sampler, batch_size=batch_size//num_gpus, **data_loader_kwargs))
if validation_set_kwargs != {}:
validation_set = dnnlib.util.construct_class_by_name(**validation_set_kwargs)
else:
validation_set = None
if rank == 0:
print()
print('Num images: ', len(training_set))
print('Image shape:', training_set.image_shape)
print('Label shape:', training_set.label_shape)
print()
if validation_set_kwargs != {}:
print()
print('Validation num images: ', len(validation_set))
print('Validation image shape: ', validation_set.image_shape)
print('Validation label shape: ', validation_set.label_shape)
# Construct networks.
if rank == 0:
print('Constructing networks...')
common_kwargs = dict(c_dim=training_set.label_dim, img_resolution=training_set.resolution, img_channels=training_set.num_channels)
G = dnnlib.util.construct_class_by_name(**G_kwargs, **common_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module
D = dnnlib.util.construct_class_by_name(**D_kwargs, **common_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module
G_ema = copy.deepcopy(G).eval()
# Resume from existing pickle.
if (resume_pkl is not None) and (rank == 0):
print(f'Resuming from "{resume_pkl}"')
with dnnlib.util.open_url(resume_pkl) as f:
resume_data = legacy.load_network_pkl(f)
for name, module in [('G', G), ('D', D), ('G_ema', G_ema)]:
misc.copy_params_and_buffers(resume_data[name], module, require_all=False)
# Print network summary tables.
if rank == 0:
z = torch.empty([batch_gpu, G.z_dim], device=device)
c = torch.empty([batch_gpu, G.c_dim], device=device)
img = misc.print_module_summary(G, [z, c])
misc.print_module_summary(D, [img, c])
# Setup augmentation.
if rank == 0:
print('Setting up augmentation...')
augment_pipe = None
ada_stats = None
if (augment_kwargs is not None) and (augment_p > 0 or ada_target is not None):
augment_pipe = dnnlib.util.construct_class_by_name(**augment_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module
augment_pipe.p.copy_(torch.as_tensor(augment_p))
if ada_target is not None:
ada_stats = training_stats.Collector(regex='Loss/signs/real')
# Distribute across GPUs.
if rank == 0:
print(f'Distributing across {num_gpus} GPUs...')
ddp_modules = dict()
for name, module in [('G_mapping', G.mapping), ('G_synthesis', G.synthesis), ('D', D), (None, G_ema), ('augment_pipe', augment_pipe)]:
if (num_gpus > 1) and (module is not None) and len(list(module.parameters())) != 0:
module.requires_grad_(True)
module = torch.nn.parallel.DistributedDataParallel(module, device_ids=[device], broadcast_buffers=False)
module.requires_grad_(False)
if name is not None:
ddp_modules[name] = module
# Setup training phases.
if rank == 0:
print('Setting up training phases...')
loss = dnnlib.util.construct_class_by_name(device=device, **ddp_modules, **loss_kwargs) # subclass of training.loss.Loss
phases = []
for name, module, opt_kwargs, reg_interval in [('G', G, G_opt_kwargs, G_reg_interval), ('D', D, D_opt_kwargs, D_reg_interval)]:
if reg_interval is None:
opt = dnnlib.util.construct_class_by_name(params=module.parameters(), **opt_kwargs) # subclass of torch.optim.Optimizer
phases += [dnnlib.EasyDict(name=name+'both', module=module, opt=opt, interval=1)]
else: # Lazy regularization.
mb_ratio = reg_interval / (reg_interval + 1)
opt_kwargs = dnnlib.EasyDict(opt_kwargs)
opt_kwargs.lr = opt_kwargs.lr * mb_ratio
opt_kwargs.betas = [beta ** mb_ratio for beta in opt_kwargs.betas]
opt = dnnlib.util.construct_class_by_name(module.parameters(), **opt_kwargs) # subclass of torch.optim.Optimizer
phases += [dnnlib.EasyDict(name=name+'main', module=module, opt=opt, interval=1)]
phases += [dnnlib.EasyDict(name=name+'reg', module=module, opt=opt, interval=reg_interval)]
for phase in phases:
phase.start_event = None
phase.end_event = None
if rank == 0:
phase.start_event = torch.cuda.Event(enable_timing=True)
phase.end_event = torch.cuda.Event(enable_timing=True)
# Export sample images.
grid_size = None
grid_z = None
grid_c = None
if rank == 0:
print('Exporting sample images...')
grid_size, images, labels = setup_snapshot_image_grid(training_set=training_set)
save_image_grid(images, os.path.join(run_dir, 'reals.png'), drange=[0,255], grid_size=grid_size)
grid_z = torch.randn([labels.shape[0], G.z_dim], device=device).split(batch_gpu)
grid_c = torch.from_numpy(labels).to(device).split(batch_gpu)
images = torch.cat([G_ema(z=z, c=c, noise_mode='const').cpu() for z, c in zip(grid_z, grid_c)]).numpy()
save_image_grid(images, os.path.join(run_dir, 'fakes_init.png'), drange=[-1,1], grid_size=grid_size)
# Initialize logs.
if rank == 0:
print('Initializing logs...')
stats_collector = training_stats.Collector(regex='.*')
stats_metrics = dict()
stats_jsonl = None
stats_tfevents = None
experiment = None
if rank == 0:
stats_jsonl = open(os.path.join(run_dir, 'stats.jsonl'), 'wt')
try:
import torch.utils.tensorboard as tensorboard
stats_tfevents = tensorboard.SummaryWriter(run_dir)
except ImportError as err:
print('Skipping tfevents export:', err)
try:
import comet_ml
except ImportError as err:
comet_ml = None
print('Skipping comet_ml export:', err)
# Train.
if rank == 0:
print(f'Training for {total_kimg} kimg...')
print()
if comet_ml is not None:
experiment = comet_ml.ExistingExperiment(
api_key=comet_api_key, previous_experiment=comet_experiment_key, auto_output_logging=False,
auto_log_co2=False, auto_metric_logging=False, auto_param_logging=False, display_summary_level=0
)
cur_nimg = 0
cur_tick = 0
tick_start_nimg = cur_nimg
tick_start_time = time.time()
maintenance_time = tick_start_time - start_time
batch_idx = 0
if progress_fn is not None:
progress_fn(0, total_kimg)
while True:
# Fetch training data.
with torch.autograd.profiler.record_function('data_fetch'):
phase_real_img, phase_real_c = next(training_set_iterator)
phase_real_img = (phase_real_img.to(device).to(torch.float32) / 127.5 - 1).split(batch_gpu)
phase_real_c = phase_real_c.to(device).split(batch_gpu)
all_gen_z = torch.randn([len(phases) * batch_size, G.z_dim], device=device)
all_gen_z = [phase_gen_z.split(batch_gpu) for phase_gen_z in all_gen_z.split(batch_size)]
all_gen_c = [training_set.get_label(np.random.randint(len(training_set))) for _ in range(len(phases) * batch_size)]
all_gen_c = torch.from_numpy(np.stack(all_gen_c)).pin_memory().to(device)
all_gen_c = [phase_gen_c.split(batch_gpu) for phase_gen_c in all_gen_c.split(batch_size)]
# Execute training phases.
for phase, phase_gen_z, phase_gen_c in zip(phases, all_gen_z, all_gen_c):
if batch_idx % phase.interval != 0:
continue
# Initialize gradient accumulation.
if phase.start_event is not None:
phase.start_event.record(torch.cuda.current_stream(device))
phase.opt.zero_grad(set_to_none=True)
phase.module.requires_grad_(True)
# Accumulate gradients over multiple rounds.
for round_idx, (real_img, real_c, gen_z, gen_c) in enumerate(zip(phase_real_img, phase_real_c, phase_gen_z, phase_gen_c)):
sync = (round_idx == batch_size // (batch_gpu * num_gpus) - 1)
gain = phase.interval
loss.accumulate_gradients(phase=phase.name, real_img=real_img, real_c=real_c, gen_z=gen_z, gen_c=gen_c, sync=sync, gain=gain)
# Update weights.
phase.module.requires_grad_(False)
with torch.autograd.profiler.record_function(phase.name + '_opt'):
for param in phase.module.parameters():
if param.grad is not None:
misc.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad)
# Clip gradients
if gradient_clipping is not None:
torch.nn.utils.clip_grad_norm_(phase.module.parameters(), gradient_clipping)
phase.opt.step()
if phase.end_event is not None:
phase.end_event.record(torch.cuda.current_stream(device))
# Update G_ema.
with torch.autograd.profiler.record_function('Gema'):
ema_nimg = ema_kimg * 1000
if ema_rampup is not None:
ema_nimg = min(ema_nimg, cur_nimg * ema_rampup)
ema_beta = 0.5 ** (batch_size / max(ema_nimg, 1e-8))
for p_ema, p in zip(G_ema.parameters(), G.parameters()):
p_ema.copy_(p.lerp(p_ema, ema_beta))
for b_ema, b in zip(G_ema.buffers(), G.buffers()):
b_ema.copy_(b)
# Update state.
cur_nimg += batch_size
batch_idx += 1
# Execute ADA heuristic.
if (ada_stats is not None) and (batch_idx % ada_interval == 0):
ada_stats.update()
adjust = np.sign(ada_stats['Loss/signs/real'] - ada_target) * (batch_size * ada_interval) / (ada_kimg * 1000)
augment_pipe.p.copy_((augment_pipe.p + adjust).max(misc.constant(0, device=device)))
# Perform maintenance tasks once per tick.
done = (cur_nimg >= total_kimg * 1000)
if (not done) and (cur_tick != 0) and (cur_nimg < tick_start_nimg + kimg_per_tick * 1000):
continue
# Print status line, accumulating the same information in stats_collector.
tick_end_time = time.time()
fields = []
fields += [f"tick {training_stats.report0('Progress/tick', cur_tick):<5d}"]
fields += [f"kimg {training_stats.report0('Progress/kimg', cur_nimg / 1e3):<8.1f}"]
fields += [f"time {dnnlib.util.format_time(training_stats.report0('Timing/total_sec', tick_end_time - start_time)):<12s}"]
fields += [f"sec/tick {training_stats.report0('Timing/sec_per_tick', tick_end_time - tick_start_time):<7.1f}"]
fields += [f"sec/kimg {training_stats.report0('Timing/sec_per_kimg', (tick_end_time - tick_start_time) / (cur_nimg - tick_start_nimg) * 1e3):<7.2f}"]
fields += [f"maintenance {training_stats.report0('Timing/maintenance_sec', maintenance_time):<6.1f}"]
fields += [f"cpumem {training_stats.report0('Resources/cpu_mem_gb', psutil.Process(os.getpid()).memory_info().rss / 2**30):<6.2f}"]
fields += [f"gpumem {training_stats.report0('Resources/peak_gpu_mem_gb', torch.cuda.max_memory_allocated(device) / 2**30):<6.2f}"]
torch.cuda.reset_peak_memory_stats()
fields += [f"augment {training_stats.report0('Progress/augment', float(augment_pipe.p.cpu()) if augment_pipe is not None else 0):.3f}"]
training_stats.report0('Timing/total_hours', (tick_end_time - start_time) / (60 * 60))
training_stats.report0('Timing/total_days', (tick_end_time - start_time) / (24 * 60 * 60))
if rank == 0:
print(' '.join(fields))
fields_for_comet = fields
# Check for abort.
if (not done) and (abort_fn is not None) and abort_fn():
done = True
if rank == 0:
print()
print('Aborting...')
# Save image snapshot.
if (rank == 0) and (image_snapshot_ticks is not None) and (done or cur_tick % image_snapshot_ticks == 0):
images = torch.cat([G_ema(z=z, c=c, noise_mode='const').cpu() for z, c in zip(grid_z, grid_c)]).numpy()
save_image_grid(images, os.path.join(run_dir, f'fakes{cur_nimg//1000:06d}.png'), drange=[-1,1], grid_size=grid_size)
# Log image and text to Comet.ml
if rank == 0 and comet_api_key and comet_ml is not None:
try:
if experiment is not None:
if (image_snapshot_ticks is not None) and (done or cur_tick % image_snapshot_ticks == 0):
experiment.log_image(image_data=os.path.join(run_dir, f'fakes{cur_nimg//1000:06d}.png'),
name=f'fakes{cur_nimg//1000:06d}', step=cur_nimg)
experiment.log_text(text=' '.join(fields_for_comet), step=cur_nimg)
except Exception as err:
print('Failed to log image to comet:', err)
# Save network snapshot.
snapshot_pkl = None
snapshot_data = None
if (network_snapshot_ticks is not None) and (done or cur_tick % network_snapshot_ticks == 0):
snapshot_data = dict(training_set_kwargs=dict(training_set_kwargs))
for name, module in [('G', G), ('D', D), ('G_ema', G_ema), ('augment_pipe', augment_pipe)]:
if module is not None:
if num_gpus > 1:
misc.check_ddp_consistency(module, ignore_regex=r'.*\.w_avg')
module = copy.deepcopy(module).eval().requires_grad_(False).cpu()
snapshot_data[name] = module
del module # conserve memory
snapshot_pkl = os.path.join(run_dir, f'network-snapshot-{cur_nimg//1000:06d}.pkl')
if rank == 0:
with open(snapshot_pkl, 'wb') as f:
pickle.dump(snapshot_data, f)
# Evaluate metrics.
if (snapshot_data is not None) and (len(metrics) > 0):
if rank == 0:
print('Evaluating metrics...')
for metric in metrics:
result_dict = metric_main.calc_metric(
metric=metric, G=snapshot_data['G_ema'], D=snapshot_data['D'], dataset_kwargs=training_set_kwargs,
validation_dataset_kwargs=validation_set_kwargs, num_gpus=num_gpus, rank=rank, device=device,
train_dataset=training_set, validation_dataset=validation_set, loss_kwargs=loss_kwargs
)
if rank == 0:
metric_main.report_metric(result_dict, run_dir=run_dir, snapshot_pkl=snapshot_pkl, cur_nimg=cur_nimg)
stats_metrics.update(result_dict.results)
# Log metrics to Comet
if experiment is not None and rank == 0:
try:
experiment.log_metrics(stats_metrics, step=cur_nimg)
except Exception as err:
print('Failed to log metrics to comet:', err)
del snapshot_data # conserve memory
# Collect statistics.
for phase in phases:
value = []
if (phase.start_event is not None) and (phase.end_event is not None):
phase.end_event.synchronize()
value = phase.start_event.elapsed_time(phase.end_event)
training_stats.report0('Timing/' + phase.name, value)
stats_collector.update()
stats_dict = stats_collector.as_dict()
# Update logs.
timestamp = time.time()
if stats_jsonl is not None:
fields = dict(stats_dict, timestamp=timestamp)
stats_jsonl.write(json.dumps(fields) + '\n')
stats_jsonl.flush()
if stats_tfevents is not None:
global_step = int(cur_nimg / 1e3)
walltime = timestamp - start_time
for name, value in stats_dict.items():
stats_tfevents.add_scalar(name, value.mean, global_step=global_step, walltime=walltime)
for name, value in stats_metrics.items():
stats_tfevents.add_scalar(f'Metrics/{name}', value, global_step=global_step, walltime=walltime)
stats_tfevents.flush()
if progress_fn is not None:
progress_fn(cur_nimg // 1000, total_kimg)
# Log losses to Comet
if experiment is not None and rank == 0:
try:
for name, value in stats_dict.items():
if name.startswith('Loss/'):
experiment.log_metric(name, value.mean, step=cur_nimg)
except Exception as err:
print('Failed to log metrics to comet:', err)
# Update state.
cur_tick += 1
tick_start_nimg = cur_nimg
tick_start_time = time.time()
maintenance_time = tick_start_time - tick_end_time
if done:
break
# Done.
if rank == 0:
print()
print('Exiting...')
run_name = config.training.resume_name
else:
run_name = "cpc" + time.strftime("-%Y-%m-%d_%H_%M_%S")
# setup logger
global_timer = timer() # global timer
logger = setup_logs(config.training.logging_dir, run_name) # setup logs
logger.info('### Experiment {} ###'.format(run_name))
logger.info('### Hyperparameter summary below ###\n {}'.format(config))
# setup of comet_ml
if has_comet:
logger.info('### Logging with comet_ml ###')
if config.comet.previous_experiment:
logger.info('===> using existing experiment: {}'.format(
config.comet.previous_experiment))
experiment = comet_ml.ExistingExperiment(
api_key=config.comet.api_key,
previous_experiment=config.comet.previous_experiment)
else:
logger.info('===> starting new experiment')
experiment = comet_ml.Experiment(api_key=config.comet.api_key,
project_name="cpc-nlp")
experiment.set_name(run_name)
experiment.log_parameters({
**config.training.to_dict(),
**config.dataset.to_dict(),
**config.cpc_model.to_dict()
})
else:
experiment = None
# define if gpu or cpu
def main(ctx, outdir, dry_run, **config_kwargs):
"""Train a GAN using the techniques described in the paper
"Training Generative Adversarial Networks with Limited Data".
Examples:
\b
# Train with custom dataset using 1 GPU.
python train.py --outdir=~/training-runs --data=~/mydataset.zip --gpus=1
\b
# Train class-conditional CIFAR-10 using 2 GPUs.
python train.py --outdir=~/training-runs --data=~/datasets/cifar10.zip \\
--gpus=2 --cfg=cifar --cond=1
\b
# Transfer learn MetFaces from FFHQ using 4 GPUs.
python train.py --outdir=~/training-runs --data=~/datasets/metfaces.zip \\
--gpus=4 --cfg=paper1024 --mirror=1 --resume=ffhq1024 --snap=10
\b
# Reproduce original StyleGAN2 config F.
python train.py --outdir=~/training-runs --data=~/datasets/ffhq.zip \\
--gpus=8 --cfg=stylegan2 --mirror=1 --aug=noaug
\b
Base configs (--cfg):
auto Automatically select reasonable defaults based on resolution
and GPU count. Good starting point for new datasets.
stylegan2 Reproduce results for StyleGAN2 config F at 1024x1024.
paper256 Reproduce results for FFHQ and LSUN Cat at 256x256.
paper512 Reproduce results for BreCaHAD and AFHQ at 512x512.
paper1024 Reproduce results for MetFaces at 1024x1024.
cifar Reproduce results for CIFAR-10 at 32x32.
\b
Transfer learning source networks (--resume):
ffhq256 FFHQ trained at 256x256 resolution.
ffhq512 FFHQ trained at 512x512 resolution.
ffhq1024 FFHQ trained at 1024x1024 resolution.
celebahq256 CelebA-HQ trained at 256x256 resolution.
lsundog256 LSUN Dog trained at 256x256 resolution.
<PATH or URL> Custom network pickle.
"""
dnnlib.util.Logger(should_flush=True)
# Setup training options.
try:
run_desc, args = setup_training_loop_kwargs(**config_kwargs)
except UserError as err:
ctx.fail(err)
# Pick output directory.
prev_run_dirs = []
if os.path.isdir(outdir):
prev_run_dirs = [
x for x in os.listdir(outdir)
if os.path.isdir(os.path.join(outdir, x))
]
prev_run_ids = [re.match(r'^\d+', x) for x in prev_run_dirs]
prev_run_ids = [int(x.group()) for x in prev_run_ids if x is not None]
cur_run_id = max(prev_run_ids, default=-1) + 1
args.run_dir = os.path.join(outdir, f'{cur_run_id:05d}-{run_desc}')
assert not os.path.exists(args.run_dir)
# Print options.
print()
print('Training options:')
print(json.dumps(args, indent=2))
print()
print(f'Output directory: {args.run_dir}')
print(f'Use comet: {bool(args.comet_api_key)}')
print(f'Training data: {args.training_set_kwargs.path}')
print(f'Training duration: {args.total_kimg} kimg')
print(f'Number of GPUs: {args.num_gpus}')
print(f'Number of images: {args.training_set_kwargs.max_size}')
print(
f'Image resolution: {args.training_set_kwargs.resolution}')
print(
f'Conditional model: {args.training_set_kwargs.use_labels}')
print(f'Dataset x-flips: {args.training_set_kwargs.xflip}')
if args.validation_set_kwargs != {}:
print(
f'Validation data: {args.validation_set_kwargs.path}')
print(
f'Number of validation images: {args.validation_set_kwargs.max_size}'
)
print()
# Dry run?
if dry_run:
print('Dry run; exiting.')
return
# Create output directory.
print('Creating output directory...')
os.makedirs(args.run_dir)
with open(os.path.join(args.run_dir, 'training_options.json'), 'wt') as f:
json.dump(args, f, indent=2)
# Setup comet experiment hyperparameters
if args.comet_api_key:
if comet_ml is not None:
try:
experiment = comet_ml.ExistingExperiment(
api_key=args.comet_api_key,
previous_experiment=args.comet_experiment_key,
auto_output_logging=False,
auto_log_co2=False,
auto_metric_logging=False,
auto_param_logging=False,
display_summary_level=0)
experiment.log_parameters(config_kwargs)
except Exception:
print('Comet logging failed')
# Launch processes.
print('Launching processes...')
torch.multiprocessing.set_start_method('spawn')
with tempfile.TemporaryDirectory() as temp_dir:
if args.num_gpus == 1:
subprocess_fn(rank=0, args=args, temp_dir=temp_dir)
else:
torch.multiprocessing.spawn(fn=subprocess_fn,
args=(args, temp_dir),
nprocs=args.num_gpus)
def predict(self, *args, **kwargs):
import comet_ml
self.experiment = comet_ml.ExistingExperiment()
def visualize(self, *args, **kwargs):
import comet_ml
self.experiment = comet_ml.ExistingExperiment()
def go():
parser = argparse.ArgumentParser(
description="Train AdmiralNet Pose Predictor")
parser.add_argument("dataset_config_file",
type=str,
help="Dataset Configuration file to load")
parser.add_argument("model_config_file",
type=str,
help="Model Configuration file to load")
parser.add_argument("output_directory",
type=str,
help="Where to put the resulting model files")
parser.add_argument(
"--context_length",
type=int,
default=None,
help="Override the context length specified in the config file")
parser.add_argument("--epochstart",
type=int,
default=1,
help="Restart training from the given epoch number")
parser.add_argument(
"--debug",
action="store_true",
help="Display images upon each iteration of the training loop")
parser.add_argument(
"--override",
action="store_true",
help="Delete output directory and replace with new data")
parser.add_argument("--tqdm",
action="store_true",
help="Display tqdm progress bar on each epoch")
parser.add_argument(
"--batch_size",
type=int,
default=None,
help="Override the order of the batch size specified in the config file"
)
parser.add_argument(
"--gpu",
type=int,
default=None,
help="Override the GPU index specified in the config file")
parser.add_argument(
"--learning_rate",
type=float,
default=None,
help="Override the learning rate specified in the config file")
parser.add_argument(
"--bezier_order",
type=int,
default=None,
help=
"Override the order of the bezier curve specified in the config file")
parser.add_argument("--weighted_loss",
action="store_true",
help="Use timewise weights on param loss")
args = parser.parse_args()
dataset_config_file = args.dataset_config_file
config_file = args.model_config_file
output_directory = os.path.join(
args.output_directory,
os.path.splitext(os.path.basename(config_file))[0])
debug = args.debug
epochstart = args.epochstart
weighted_loss = args.weighted_loss
with open(config_file) as f:
config = yaml.load(f, Loader=yaml.SafeLoader)
with open(dataset_config_file) as f:
dataset_config = yaml.load(f, Loader=yaml.SafeLoader)
image_size = dataset_config["image_size"]
input_channels = config["input_channels"]
if args.context_length is not None:
context_length = args.context_length
config["context_length"] = context_length
output_directory += "_context%d" % (context_length)
else:
context_length = config["context_length"]
if args.bezier_order is not None:
bezier_order = args.bezier_order
config["bezier_order"] = bezier_order
output_directory += "_bezier_order%d" % (bezier_order)
else:
bezier_order = config["bezier_order"]
#num_recurrent_layers = config["num_recurrent_layers"]
if args.gpu is not None:
gpu = args.gpu
config["gpu"] = gpu
else:
gpu = config["gpu"]
torch.cuda.set_device(gpu)
if args.batch_size is not None:
batch_size = args.batch_size
config["batch_size"] = batch_size
else:
batch_size = config["batch_size"]
if args.learning_rate is not None:
learning_rate = args.learning_rate
config["learning_rate"] = learning_rate
else:
learning_rate = config["learning_rate"]
momentum = config["momentum"]
num_epochs = config["num_epochs"]
num_workers = config["num_workers"]
use_float = config["use_float"]
loss_weights = config["loss_weights"]
hidden_dim = config["hidden_dimension"]
num_recurrent_layers = config.get("num_recurrent_layers", 1)
config["hostname"] = socket.gethostname()
print("Using config:\n%s" % (str(config)))
net = nn_models.Models.AdmiralNetCombinedBezierPredictor(
context_length=context_length,
hidden_dim=hidden_dim,
num_recurrent_layers=num_recurrent_layers,
params_per_dimension=bezier_order + 1)
print("net:\n%s" % (str(net)))
ppd = net.pos_predictor.params_per_dimension
numones = int(ppd / 2)
if weighted_loss:
timewise_weights = torch.from_numpy(
np.hstack((np.ones(numones), np.linspace(1, 3, ppd - numones))))
else:
timewise_weights = None
params_loss = loss_functions.SquaredLpNormLoss(
timewise_weights=timewise_weights)
kinematic_loss = loss_functions.SquaredLpNormLoss()
if use_float:
print("casting stuff to float")
net = net.float()
params_loss = params_loss.float()
kinematic_loss = kinematic_loss.float()
else:
print("casting stuff to double")
net = net.double()
params_loss = params_loss.double()
kinematic_loss = kinematic_loss.float()
if gpu >= 0:
print("moving stuff to GPU")
net = net.cuda()
params_loss = params_loss.cuda()
kinematic_loss = kinematic_loss.cuda()
optimizer = optim.SGD(net.parameters(),
lr=learning_rate,
momentum=momentum,
dampening=0.000,
nesterov=True)
netpostfix = "epoch_%d_params.pt"
optimizerpostfix = "epoch_%d_optimizer.pt"
if epochstart > 1:
net.load_state_dict(
torch.load(os.path.join(output_directory,
netpostfix % (epochstart)),
map_location=next(net.parameters()).device))
optimizer.load_state_dict(
torch.load(os.path.join(output_directory,
optimizerpostfix % (epochstart)),
map_location=next(net.parameters()).device))
experiment_config = yaml.load(open(
os.path.join(output_directory, "experiment_config.yaml"), "r"),
Loader=yaml.SafeLoader)
experiment = comet_ml.ExistingExperiment(
workspace="electric-turtle",
project_name="deepracingbezierpredictor",
previous_experiment=experiment_config["experiment_key"])
else:
if (not args.override) and os.path.isdir(output_directory):
s = ""
while (not (s == "y" or s == "n")):
s = input(
"Directory " + output_directory +
" already exists. Overwrite it with new data? [y\\n]\n")
if s == "n":
print("Thanks for playing!")
exit(0)
shutil.rmtree(output_directory)
elif os.path.isdir(output_directory):
shutil.rmtree(output_directory)
os.makedirs(output_directory, exist_ok=True)
experiment = comet_ml.Experiment(
workspace="electric-turtle",
project_name="deepracingbezierpredictor")
experiment.log_parameters(config)
experiment.log_parameters(dataset_config)
experiment.add_tag("bezierpredictor")
experiment_config = {"experiment_key": experiment.get_key()}
yaml.dump(experiment_config,
stream=open(
os.path.join(output_directory, "experiment_config.yaml"),
"w"),
Dumper=yaml.SafeDumper)
if num_workers == 0:
max_spare_txns = 50
else:
max_spare_txns = num_workers
#image_wrapper = deepracing.backend.ImageFolderWrapper(os.path.dirname(image_db))
dsets = []
use_optflow = True
for dataset in dataset_config["datasets"]:
print("Parsing database config: %s" % (str(dataset)))
label_folder = dataset["label_folder"]
key_file = dataset["key_file"]
image_folder = dataset["image_folder"]
apply_color_jitter = dataset.get("apply_color_jitter", False)
erasing_probability = dataset.get("erasing_probability", 0.0)
label_wrapper = deepracing.backend.PoseSequenceLabelLMDBWrapper()
label_wrapper.readDatabase(os.path.join(label_folder, "lmdb"),
max_spare_txns=max_spare_txns)
image_mapsize = float(np.prod(image_size) * 3 + 12) * float(
len(label_wrapper.getKeys())) * 1.1
image_wrapper = deepracing.backend.ImageLMDBWrapper(
direct_caching=False)
image_wrapper.readDatabase(os.path.join(image_folder, "image_lmdb"),
max_spare_txns=max_spare_txns,
mapsize=image_mapsize)
curent_dset = data_loading.proto_datasets.PoseSequenceDataset(image_wrapper, label_wrapper, key_file, context_length,\
image_size = image_size, return_optflow=use_optflow, apply_color_jitter=apply_color_jitter, erasing_probability=erasing_probability)
dsets.append(curent_dset)
print("\n")
if len(dsets) == 1:
dset = dsets[0]
else:
dset = torch.utils.data.ConcatDataset(dsets)
dataloader = data_utils.DataLoader(dset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=gpu >= 0)
print("Dataloader of of length %d" % (len(dataloader)))
yaml.dump(dataset_config,
stream=open(
os.path.join(output_directory, "dataset_config.yaml"), "w"),
Dumper=yaml.SafeDumper)
yaml.dump(config,
stream=open(os.path.join(output_directory, "config.yaml"), "w"),
Dumper=yaml.SafeDumper)
if (epochstart == 1):
i = 0
else:
i = epochstart
with experiment.train():
while i < num_epochs:
time.sleep(2.0)
postfix = i + 1
print("Running Epoch Number %d" % (postfix))
#dset.clearReaders()
try:
tick = time.time()
run_epoch(experiment,
net,
optimizer,
dataloader,
gpu,
params_loss,
kinematic_loss,
loss_weights,
debug=debug,
use_tqdm=args.tqdm)
tock = time.time()
print("Finished epoch %d in %f seconds." %
(postfix, tock - tick))
experiment.log_epoch_end(postfix)
except Exception as e:
print("Restarting epoch %d because %s" % (postfix, str(e)))
modelin = os.path.join(output_directory,
netpostfix % (postfix - 1))
optimizerin = os.path.join(output_directory,
optimizerpostfix % (postfix - 1))
net.load_state_dict(torch.load(modelin))
optimizer.load_state_dict(torch.load(optimizerin))
continue
modelout = os.path.join(output_directory, netpostfix % (postfix))
torch.save(net.state_dict(), modelout)
with open(modelout, 'rb') as modelfile:
experiment.log_asset(modelfile,
file_name=netpostfix % (postfix))
optimizerout = os.path.join(output_directory,
optimizerpostfix % (postfix))
torch.save(optimizer.state_dict(), optimizerout)
with open(optimizerout, 'rb') as optimizerfile:
experiment.log_asset(optimizerfile,
file_name=optimizerpostfix % (postfix))
i = i + 1
def go():
parser = argparse.ArgumentParser(description="Train AdmiralNet Pose Predictor")
parser.add_argument("model_directory", type=str, help="Where to put the resulting model files")
parser.add_argument("epochstart", type=int, default=None, help="Restart training from the given epoch number")
parser.add_argument("--context_length", type=int, default=None, help="Override the context length specified in the config file")
parser.add_argument("--epochstart", type=int, default=None, help="Restart training from the given epoch number")
parser.add_argument("--debug", action="store_true", help="Display images upon each iteration of the training loop")
parser.add_argument("--override", action="store_true", help="Delete output directory and replace with new data")
parser.add_argument("--tqdm", action="store_true", help="Display tqdm progress bar on each epoch")
parser.add_argument("--batch_size", type=int, default=None, help="Override the order of the batch size specified in the config file")
parser.add_argument("--gpu", type=int, default=None, help="Override the GPU index specified in the config file")
parser.add_argument("--learning_rate", type=float, default=None, help="Override the learning rate specified in the config file")
parser.add_argument("--momentum", type=float, default=None, help="Override the momentum specified in the config file")
parser.add_argument("--dampening", type=float, default=None, help="Override the dampening specified in the config file")
parser.add_argument("--nesterov", action="store_true", help="Override the nesterov specified in the config file")
parser.add_argument("--bezier_order", type=int, default=None, help="Override the order of the bezier curve specified in the config file")
parser.add_argument("--optimizer", type=str, default="SGD", help="Optimizer to use")
parser.add_argument("--velocity_loss", type=float, default=None, help="Override velocity loss weight in config file")
parser.add_argument("--position_loss", type=float, default=None, help="Override position loss weight in config file")
parser.add_argument("--control_point_loss", type=float, default=None, help="Override control point loss weight in config file")
parser.add_argument("--fix_first_point",type=bool,default=False, help="Override fix_first_point in the config file")
args = parser.parse_args()
output_directory = args.model_directory
epochstart = args.epochstart
dataset_config_file = os.path.join(output_directory,"dataset_config.yaml")
with open(dataset_config_file) as f:
dataset_config = yaml.load(f, Loader = yaml.SafeLoader)
config_file = os.path.join(output_directory,"model_config.yaml")
with open(config_file) as f:
config = yaml.load(f, Loader = yaml.SafeLoader)
experiment_config_file = os.path.join(output_directory,"experiment_config.yaml")
with open(experiment_config_file) as f:
experiment_config = yaml.load(f, Loader = yaml.SafeLoader)
experiment_id = experiment_config["experiment_key"]
print(dataset_config)
image_size = dataset_config["image_size"]
input_channels = config["input_channels"]
if args.context_length is not None:
context_length = args.context_length
config["context_length"] = context_length
else:
context_length = config["context_length"]
if args.bezier_order is not None:
bezier_order = args.bezier_order
config["bezier_order"] = bezier_order
else:
bezier_order = config["bezier_order"]
#num_recurrent_layers = config["num_recurrent_layers"]
if args.gpu is not None:
gpu = args.gpu
config["gpu"] = gpu
else:
gpu = config["gpu"]
torch.cuda.set_device(gpu)
if args.batch_size is not None:
batch_size = args.batch_size
config["batch_size"] = batch_size
else:
batch_size = config["batch_size"]
if args.learning_rate is not None:
learning_rate = args.learning_rate
config["learning_rate"] = learning_rate
else:
learning_rate = config["learning_rate"]
if args.momentum is not None:
momentum = args.momentum
config["momentum"] = momentum
else:
momentum = config["momentum"]
if args.dampening is not None:
dampening = args.dampening
config["dampening"] = dampening
else:
dampening = config["dampening"]
if args.nesterov:
nesterov = True
config["nesterov"] = nesterov
else:
nesterov = config["nesterov"]
if args.fix_first_point:
fix_first_point = True
config["fix_first_point"] = fix_first_point
else:
fix_first_point = config["fix_first_point"]
loss_weights = config["loss_weights"]
if args.control_point_loss is not None:
loss_weights[0] = args.control_point_loss
config["loss_weights"] = loss_weights
if args.position_loss is not None:
loss_weights[1] = args.position_loss
config["loss_weights"] = loss_weights
if args.velocity_loss is not None:
loss_weights[2] = args.velocity_loss
config["loss_weights"] = loss_weights
num_epochs = config["num_epochs"]
num_workers = config["num_workers"]
use_float = config["use_float"]
hidden_dim = config["hidden_dimension"]
loss_reduction = config["loss_reduction"]
use_3dconv = config["use_3dconv"]
num_recurrent_layers = config.get("num_recurrent_layers",1)
config["hostname"] = socket.gethostname()
print("Using config:\n%s" % (str(config)))
net = deepracing_models.nn_models.Models.AdmiralNetCurvePredictor( context_length = context_length , input_channels=input_channels, hidden_dim = hidden_dim, num_recurrent_layers=num_recurrent_layers, params_per_dimension=bezier_order + 1 - int(fix_first_point), use_3dconv = use_3dconv)
print("net:\n%s" % (str(net)))
ppd = net.params_per_dimension
numones = int(ppd/2)
weighted_loss = config.get("weighted_loss", False)
if weighted_loss:
timewise_weights = torch.from_numpy( np.hstack( ( np.ones(numones), np.linspace(1,3, ppd - numones ) ) ) )
else:
timewise_weights = None
params_loss = deepracing_models.nn_models.LossFunctions.SquaredLpNormLoss(time_reduction=loss_reduction)
kinematic_loss = deepracing_models.nn_models.LossFunctions.SquaredLpNormLoss(time_reduction=loss_reduction)
weightfilename = os.path.join(output_directory,"epoch_%d_params.pt" %(epochstart,))
optimizerfilename = os.path.join(output_directory,"epoch_%d_optimizer.pt" %(epochstart,))
#get network weights
with open(weightfilename,"rb") as f:
net.load_state_dict(torch.load(f, map_location=torch.device("cpu")))
#get optimizer weights
if gpu>=0:
print("moving stuff to GPU")
net = net.cuda(gpu)
params_loss = params_loss.cuda(gpu)
kinematic_loss = kinematic_loss.cuda(gpu)
optimizer = config["optimizer"]
if optimizer=="Adam":
optimizer = optim.Adam(net.parameters(), lr = learning_rate, betas=(0.9, 0.9))
elif optimizer=="RMSprop":
optimizer = optim.RMSprop(net.parameters(), lr = learning_rate, momentum = momentum)
elif optimizer=="ASGD":
optimizer = optim.ASGD(net.parameters(), lr = learning_rate)
elif optimizer=="SGD":
nesterov_ = momentum>0.0 and nesterov
dampening_=dampening*float(not nesterov_)
optimizer = optim.SGD(net.parameters(), lr = learning_rate, momentum = momentum, dampening=dampening_, nesterov=nesterov_)
else:
raise ValueError("Uknown optimizer " + optimizer)
with open(optimizerfilename,"rb") as f:
optimizer.load_state_dict(torch.load(f, map_location=torch.device("cpu")))
if use_float:
print("casting stuff to float")
net = net.float()
params_loss = params_loss.float()
kinematic_loss = kinematic_loss.float()
else:
print("casting stuff to double")
net = net.double()
params_loss = params_loss.double()
kinematic_loss = kinematic_loss.float()
if num_workers == 0:
max_spare_txns = 50
else:
max_spare_txns = num_workers
#image_wrapper = deepracing.backend.ImageFolderWrapper(os.path.dirname(image_db))
dsets=[]
use_optflow = net.input_channels==5
dsetfolders = []
apply_color_jitter = dataset_config.get("apply_color_jitter",False)
for dataset in dataset_config["datasets"]:
print("Parsing database config: %s" %(str(dataset)))
root_folder = dataset["root_folder"]
dsetfolders.append(root_folder)
label_folder = os.path.join(root_folder,"pose_sequence_labels")
image_folder = os.path.join(root_folder,"images")
key_file = os.path.join(root_folder,"goodkeys.txt")
erasing_probability = dataset.get("erasing_probability",0.0)
label_wrapper = deepracing.backend.PoseSequenceLabelLMDBWrapper()
label_wrapper.readDatabase(os.path.join(label_folder,"lmdb") )
image_mapsize = float(np.prod(image_size)*3+12)*float(len(label_wrapper.getKeys()))*1.1
image_wrapper = deepracing.backend.ImageLMDBWrapper(direct_caching=False)
image_wrapper.readDatabase(os.path.join(image_folder,"image_lmdb"), mapsize=image_mapsize )
curent_dset = PD.PoseSequenceDataset(image_wrapper, label_wrapper, key_file, context_length,\
image_size = image_size, return_optflow=use_optflow, apply_color_jitter=apply_color_jitter, erasing_probability=erasing_probability)
dsets.append(curent_dset)
print("\n")
if len(dsets)==1:
dset = dsets[0]
else:
dset = torch.utils.data.ConcatDataset(dsets)
dataloader = data_utils.DataLoader(dset, batch_size=batch_size,
shuffle=True, num_workers=num_workers, pin_memory=gpu>=0)
print("Dataloader of of length %d" %(len(dataloader)))
netpostfix = "epoch_%d_params.pt"
optimizerpostfix = "epoch_%d_optimizer.pt"
experiment = comet_ml.ExistingExperiment(workspace="electric-turtle", project_name="deepracingbezierpredictor", \
previous_experiment=experiment_id)
i = epochstart
with experiment.train():
while i < num_epochs:
time.sleep(2.0)
postfix = i + 1
print("Running Epoch Number %d" %(postfix))
#dset.clearReaders()
try:
tick = time.time()
run_epoch(experiment, net, fix_first_point, optimizer, dataloader, gpu, params_loss, kinematic_loss, loss_weights, postfix, debug=False, use_tqdm=args.tqdm )
tock = time.time()
print("Finished epoch %d in %f seconds." % ( postfix , tock-tick ) )
experiment.log_epoch_end(postfix)
except Exception as e:
if isinstance(e, FileExistsError):
print(e)
exit(-1)
print("Restarting epoch %d because %s"%(postfix, str(e)))
modelin = os.path.join(output_directory, netpostfix %(postfix-1))
optimizerin = os.path.join(output_directory,optimizerpostfix %(postfix-1))
net.load_state_dict(torch.load(modelin))
optimizer.load_state_dict(torch.load(optimizerin))
continue
modelout = os.path.join(output_directory,netpostfix %(postfix))
torch.save(net.state_dict(), modelout)
with open(modelout,'rb') as modelfile:
experiment.log_asset(modelfile,file_name=netpostfix %(postfix))
optimizerout = os.path.join(output_directory,optimizerpostfix %(postfix))
torch.save(optimizer.state_dict(), optimizerout)
with open(optimizerout,'rb') as optimizerfile:
experiment.log_asset(optimizerfile,file_name=optimizerpostfix %(postfix))
i = i + 1
def predict(self, *args, **kwargs):
import comet_ml
self.experiment = comet_ml.ExistingExperiment()
self._log_html(" ".join(args))