def main(args):
global IMG_SIZE
IMG_SIZE = (args.img_size, args.img_size)
model = ResnetEncoder(arch=args.arch, pooling=args.pooling)
model = model.eval()
model, device = UtilsFactory.prepare_model(model)
images_df = pd.read_csv(args.in_csv)
images_df = images_df.reset_index().drop("index", axis=1)
images_df = list(images_df.to_dict("index").values())
open_fn = ImageReader(row_key=args.img_col,
dict_key="image",
datapath=args.datapath)
dataloader = UtilsFactory.create_loader(images_df,
open_fn,
batch_size=args.batch_size,
workers=args.n_workers,
dict_transform=dict_transformer)
features = []
dataloader = tqdm(dataloader) if args.verbose else dataloader
with torch.no_grad():
for batch in dataloader:
features_ = model(batch["image"].to(device))
features_ = features_.cpu().detach().numpy()
features.append(features_)
features = np.concatenate(features, axis=0)
np.save(args.out_npy, features)
def _init(self,
critics,
min_action,
max_action,
action_noise_std=0.2,
action_noise_clip=0.5,
**kwargs):
super()._init(**kwargs)
self.min_action = min_action
self.max_action = max_action
self.action_noise_std = action_noise_std
self.action_noise_clip = action_noise_clip
critics = [x.to(self._device) for x in critics]
critics_optimizer = [
UtilsFactory.create_optimizer(x, **self.critic_optimizer_params)
for x in critics
]
critics_scheduler = [
UtilsFactory.create_scheduler(x, **self.critic_scheduler_params)
for x in critics_optimizer
]
target_critics = [copy.deepcopy(x).to(self._device) for x in critics]
self.critics = [self.critic] + critics
self.critics_optimizer = [self.critic_optimizer] + critics_optimizer
self.critics_scheduler = [self.critic_scheduler] + critics_scheduler
self.target_critics = [self.target_critic] + target_critics
def main(args):
model = models.__dict__[args.arch](pretrained=True)
model = model.eval()
model, device = UtilsFactory.prepare_model(model)
labels = json.loads(open(args.labels).read())
i2k = Images2Keywords(model, args.n_keywords, labels)
images_df = pd.read_csv(args.in_csv)
images_df = images_df.reset_index().drop("index", axis=1)
images_df = list(images_df.to_dict("index").values())
open_fn = ImageReader(row_key=args.img_col,
dict_key="image",
datapath=args.datapath)
dataloader = UtilsFactory.create_loader(images_df,
open_fn,
batch_size=args.batch_size,
workers=args.n_workers,
dict_transform=dict_transformer)
keywords = []
dataloader = tqdm(dataloader) if args.verbose else dataloader
with torch.no_grad():
for batch in dataloader:
keywords_batch = i2k(batch["image"].to(device))
keywords += keywords_batch
input_csv = pd.read_csv(args.in_csv)
input_csv[args.keywords_col] = keywords
input_csv.to_csv(args.out_csv, index=False)
def to_batch_metrics(*, state, metric_key):
metric = state.get_key(metric_key)
if isinstance(metric, dict):
for key, value in metric.items():
state.batch_metrics[f"{metric_key}_{key}"] = \
UtilsFactory.get_val_from_metric(value)
else:
state.batch_metrics[f"{metric_key}"] = \
UtilsFactory.get_val_from_metric(metric)
def _init(self, critic_v, **kwargs):
super()._init(**kwargs)
self.critic_v = critic_v.to(self._device)
self.critic_v_optimizer = UtilsFactory.create_optimizer(
self.critic_v, **self.critic_optimizer_params)
self.critic_v_scheduler = UtilsFactory.create_scheduler(
self.critic_v_optimizer, **self.critic_scheduler_params)
self.target_critic_v = copy.deepcopy(critic_v).to(self._device)
self.num_atoms = self.critic.n_atoms
tau_min = 1 / (2 * self.num_atoms)
tau_max = 1 - tau_min
tau = torch.linspace(start=tau_min, end=tau_max, steps=self.num_atoms)
self.tau = self.to_tensor(tau)
def _init(self):
"""
Inner method for children's classes for model specific initialization.
As baseline, checks device support and puts model on it.
:return:
"""
self.model, self.device = UtilsFactory.prepare_model(self.model)
def main(args):
images_df = pd.read_csv(args.in_csv)
images_df = images_df.reset_index().drop("index", axis=1)
images_df = list(images_df.to_dict("index").values())
if args.fasttext_model is not None:
encode_fn = create_fasttext_encode_fn(args.fasttext_model,
normalize=args.normalize)
elif args.w2v_model is not None:
encode_fn = create_gensim_encode_fn(args.w2v_model,
sep=args.txt_sep,
normalize=args.normalize)
else:
raise NotImplementedError
open_fn = TextReader(row_key=args.txt_col,
dict_key="txt",
encode_fn=encode_fn)
dataloader = UtilsFactory.create_loader(images_df,
open_fn,
batch_size=args.batch_size,
workers=args.n_workers)
features = []
dataloader = tqdm(dataloader) if args.verbose else dataloader
for batch in dataloader:
features_ = batch["txt"]
features.append(features_)
features = np.concatenate(features, axis=0)
np.save(args.out_npy, features)
def save(self):
if self.epoch % self.save_period == 0:
checkpoint = self.algorithm.prepare_checkpoint()
checkpoint["epoch"] = self.epoch
filename = UtilsFactory.save_checkpoint(logdir=self.logdir,
checkpoint=checkpoint,
suffix=str(self.epoch))
print("Checkpoint saved to: %s" % filename)
def on_batch_end(self, state):
lm = state.loader_mode
if lm not in self.loggers:
self.loggers[lm] = UtilsFactory.create_tflogger(logdir=self.logdir,
name=lm)
for key, value in state.batch_metrics.items():
self.loggers[lm].add_scalar(key, value, state.step)
def load_checkpoint(*, filename, state):
if os.path.isfile(filename):
print("=> loading checkpoint \"{}\"".format(filename))
checkpoint = UtilsFactory.load_checkpoint(filename)
state.epoch = checkpoint["epoch"]
state.best_metrics = checkpoint["best_metrics"]
UtilsFactory.unpack_checkpoint(checkpoint,
model=state.model,
criterion=state._criterion,
optimizer=state._optimizer,
scheduler=state._scheduler)
print("loaded checkpoint \"{}\" (epoch {})".format(
filename, checkpoint["epoch"]))
else:
raise Exception("no checkpoint found at \"{}\"".format(filename))
def _init(self,
critic_q1,
critic_q2,
reward_scale=1.0,
use_regularization=False,
mu_and_sigma_reg=1e-3,
policy_grad_estimator="reparametrization_trick",
**kwargs):
"""
Parameters
----------
reward_scale: float,
THE MOST IMPORTANT HYPERPARAMETER which controls the ratio
between maximizing rewards and acting as randomly as possible
use_regularization: bool,
whether to use l2 regularization on policy network outputs,
regularization can not be used with RealNVPActor
mu_and_sigma_reg: float,
coefficient for l2 regularization on mu and log_sigma
policy_grad_estimator: str,
"reinforce": may be used with arbitrary explicit policy
"reparametrization_trick": may be used with reparametrizable
policy, e.g. Gaussian, normalizing flow (Real NVP).
"""
super()._init(**kwargs)
self.critic_q1 = critic_q1.to(self._device)
self.critic_q2 = critic_q2.to(self._device)
self.critic_q1_optimizer = UtilsFactory.create_optimizer(
self.critic_q1, **self.critic_optimizer_params)
self.critic_q2_optimizer = UtilsFactory.create_optimizer(
self.critic_q2, **self.critic_optimizer_params)
self.critic_q1_scheduler = UtilsFactory.create_scheduler(
self.critic_q1_optimizer, **self.critic_scheduler_params)
self.critic_q2_scheduler = UtilsFactory.create_scheduler(
self.critic_q2_optimizer, **self.critic_scheduler_params)
self.reward_scale = reward_scale
self.use_regularization = use_regularization
self.mu_sigma_reg = mu_and_sigma_reg
self.policy_grad_estimator = policy_grad_estimator
def _init(self,
critic2,
min_action,
max_action,
action_noise_std=0.2,
action_noise_clip=0.5,
**kwargs):
super()._init(**kwargs)
self.min_action = min_action
self.max_action = max_action
self.action_noise_std = action_noise_std
self.action_noise_clip = action_noise_clip
self.critic2 = critic2.to(self._device)
self.critic2_optimizer = UtilsFactory.create_optimizer(
self.critic2, **self.critic_optimizer_params)
self.critic2_scheduler = UtilsFactory.create_scheduler(
self.critic_optimizer, **self.critic_scheduler_params)
self.target_critic2 = copy.deepcopy(critic2).to(self._device)
def load_actor_weights(self):
if self.resume is not None:
checkpoint = UtilsFactory.load_checkpoint(self.resume)
weights = checkpoint[f"actor_state_dict"]
self.actor.load_state_dict(weights)
elif self.redis_server is not None:
weights = deserialize(
self.redis_server.get(f"{self.redis_prefix}_actor_weights"))
weights = {k: self.to_tensor(v) for k, v in weights.items()}
self.actor.load_state_dict(weights)
else:
raise NotImplementedError
self.actor.eval()
def save_checkpoint(self, logdir, checkpoint, is_best, save_n_best=5):
suffix = f"{checkpoint['stage']}.{checkpoint['epoch']}"
filepath = UtilsFactory.save_checkpoint(logdir=logdir,
checkpoint=checkpoint,
is_best=is_best,
suffix=suffix)
self.top_best_metrics.append(
(filepath, checkpoint["valid_metrics"][self.main_metric]))
self.top_best_metrics = sorted(self.top_best_metrics,
key=lambda x: x[1],
reverse=not self.minimize)
if len(self.top_best_metrics) > save_n_best:
last_item = self.top_best_metrics.pop(-1)
last_filepath = last_item[0]
os.remove(last_filepath)
def load_checkpoint(self, filepath, load_optimizer=True):
checkpoint = UtilsFactory.load_checkpoint(filepath)
for key in ["actor", "critic"]:
value_l = getattr(self, key, None)
if value_l is not None:
value_r = checkpoint[f"{key}_state_dict"]
value_l.load_state_dict(value_r)
if load_optimizer:
for key2 in ["optimizer", "scheduler"]:
key2 = f"{key}_{key2}"
value_l = getattr(self, key2, None)
if value_l is not None:
value_r = checkpoint[f"{key2}_state_dict"]
value_l.load_state_dict(value_r)
def load_checkpoint(self, filepath, load_optimizer=True):
super().load_checkpoint(filepath, load_optimizer)
checkpoint = UtilsFactory.load_checkpoint(filepath)
key = "critics"
for i in range(len(self.critics)):
value_l = getattr(self, key, None)
value_l = value_l[i] if value_l is not None else None
if value_l is not None:
value_r = checkpoint[f"{key}{i}_state_dict"]
value_l.load_state_dict(value_r)
if load_optimizer:
for key2 in ["optimizer", "scheduler"]:
key2 = f"{key}_{key2}"
value_l = getattr(self, key2, None)
if value_l is not None:
value_r = checkpoint[f"{key2}_state_dict"]
value_l.load_state_dict(value_r)
def __init__(self,
grad_clip_params: Dict = None,
fp16_grad_scale: float = 128.0,
accumulation_steps: int = 1,
optimizer_key: str = None,
loss_key: str = None):
"""
@TODO: docs
"""
grad_clip_params = grad_clip_params or {}
self.grad_clip_fn = UtilsFactory.create_grad_clip_fn(
**grad_clip_params)
self.fp16 = False
self.fp16_grad_scale = fp16_grad_scale
self.accumulation_steps = accumulation_steps
self.optimizer_key = optimizer_key
self.loss_key = loss_key
self.optimizer_wd = 0
self.accumulation_counter = 0
def save_checkpoint(
self, logdir, checkpoint, is_best,
save_n_best=5, main_metric="loss_main", minimize_metric=True):
suffix = f"{checkpoint['stage']}.{checkpoint['epoch']}"
filepath = UtilsFactory.save_checkpoint(
logdir=logdir,
checkpoint=checkpoint,
is_best=is_best,
suffix=suffix)
checkpoint_metric = checkpoint["valid_metrics"].get(main_metric, None)
checkpoint_metric = checkpoint_metric or checkpoint.get("epoch", -1)
self.top_best_metrics.append((filepath, checkpoint_metric))
self.top_best_metrics = sorted(
self.top_best_metrics,
key=lambda x: x[1],
reverse=not minimize_metric)
if len(self.top_best_metrics) > save_n_best:
last_item = self.top_best_metrics.pop(-1)
last_filepath = last_item[0]
os.remove(last_filepath)
def on_epoch_end_pre(state):
if state.mode == "infer":
return
best_metrics, valid_metrics, is_best = \
UtilsFactory.process_epoch_metrics(
state.epoch_metrics,
state.best_metrics,
valid_loader=state.valid_loader,
main_metric=state.main_metric,
minimize=state.minimize_metric)
valid_metrics = {
key: value
for key, value in valid_metrics.items()
if isinstance(value, float)
}
state.best_metrics = {
key: value
for key, value in best_metrics.items() if isinstance(value, float)
}
state.valid_metrics = valid_metrics
state.is_best_epoch = is_best
def prepare_model(config):
return UtilsFactory.create_model(config=config,
available_networks=NETWORKS)
def pack_checkpoint(self, **kwargs):
return UtilsFactory.pack_checkpoint(**kwargs)
def on_loader_end_pre(state):
lm = state.loader_mode
state.epoch_metrics[lm] = {
key: UtilsFactory.get_val_from_metric(value)
for key, value in state.epoch_metrics[lm].items()
}
def on_loader_start(self, state):
lm = state.loader_mode
if lm not in self.loggers:
self.loggers[lm] = UtilsFactory.create_tflogger(logdir=self.logdir,
name=lm)
def __init__(self,
actor,
critic,
gamma,
n_step,
actor_optimizer_params,
critic_optimizer_params,
actor_grad_clip=None,
critic_grad_clip=None,
actor_loss_params=None,
critic_loss_params=None,
actor_scheduler_params=None,
critic_scheduler_params=None,
resume=None,
load_optimizer=True,
actor_tau=1.0,
critic_tau=1.0,
**kwargs):
self._device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self.actor = actor.to(self._device)
self.critic = critic.to(self._device)
self.target_actor = copy.deepcopy(actor).to(self._device)
self.target_critic = copy.deepcopy(critic).to(self._device)
self.actor_optimizer = UtilsFactory.create_optimizer(
self.actor, **actor_optimizer_params)
self.critic_optimizer = UtilsFactory.create_optimizer(
self.critic, **critic_optimizer_params)
self.actor_optimizer_params = actor_optimizer_params
self.critic_optimizer_params = critic_optimizer_params
self.actor_scheduler = UtilsFactory.create_scheduler(
self.actor_optimizer, **actor_scheduler_params)
self.critic_scheduler = UtilsFactory.create_scheduler(
self.critic_optimizer, **critic_scheduler_params)
self.actor_scheduler_params = actor_scheduler_params
self.critic_scheduler_params = critic_scheduler_params
self.n_step = n_step
self.gamma = gamma
self.actor_grad_clip_params = None
self.critic_grad_clip_params = None
if actor_grad_clip is not None:
# clip_grad_value_ or clip_grad_norm_
func = actor_grad_clip.pop("func", "clip_grad_value_")
func = torch.nn.utils.__dict__[func]
actor_grad_clip_params = copy.deepcopy(actor_grad_clip)
actor_grad_clip = lambda parameters: func(parameters, **
actor_grad_clip_params)
self.actor_grad_clip_params = actor_grad_clip_params
if critic_grad_clip is not None:
# clip_grad_value_ or clip_grad_norm_
func = critic_grad_clip.pop("func", "clip_grad_value_")
func = torch.nn.utils.__dict__[func]
critic_grad_clip_params = copy.deepcopy(critic_grad_clip)
critic_grad_clip = lambda parameters: func(
parameters, **critic_grad_clip_params)
self.critic_grad_clip_params = critic_grad_clip_params
self.actor_grad_clip = actor_grad_clip
self.critic_grad_clip = critic_grad_clip
self.actor_tau = actor_tau
self.critic_tau = critic_tau
self.actor_criterion = UtilsFactory.create_criterion(
**(actor_loss_params or {}))
self.critic_criterion = UtilsFactory.create_criterion(
**(critic_loss_params or {}))
self.actor_loss_params = actor_loss_params
self.critic_loss_params = critic_loss_params
self._init(**kwargs)
if resume is not None:
self.load_checkpoint(resume, load_optimizer=load_optimizer)
augment_fn=lambda x: \
torch.from_numpy(x.copy().astype(np.float32) / 256.).unsqueeze_(0).float()),
Augmentor(
dict_key="targets",
augment_fn=lambda x: \
torch.from_numpy(x.copy()).unsqueeze_(0).float()),
]
train_data_transform = transforms.Compose(augmentations + transformations)
valid_data_transform = transforms.Compose(transformations)
open_fn = lambda x: {"features": x[0], "targets": x[1]}
train_loader = UtilsFactory.create_loader(train_data,
open_fn=open_fn,
dict_transform=train_data_transform,
batch_size=bs,
workers=n_workers,
shuffle=True)
valid_loader = UtilsFactory.create_loader(valid_data,
open_fn=open_fn,
dict_transform=valid_data_transform,
batch_size=bs,
workers=n_workers,
shuffle=False)
loaders = collections.OrderedDict()
loaders["train"] = train_loader
loaders["valid"] = valid_loader
# Model & Criterion
def train(self,
*,
datasource: AbstractDataSource,
args: Namespace,
stages_config: Dict[str, Dict] = None,
verbose: bool = False):
"""
Main method for training DL models.
:param datasource: AbstractDataSource instance
:param args: console args
:param stages_config: config
:param verbose: verbose flag
"""
stages_state_params = stages_config.pop("state_params", {})
stages_data_params = stages_config.pop("data_params", {})
stages_callbacks_params = stages_config.pop("callbacks_params", {})
stages_criterion_params = stages_config.pop("criterion_params", {})
stages_optimizer_params = stages_config.pop("optimizer_params", {})
stages_scheduler_params = stages_config.pop("scheduler_params", {})
loaders = None
for stage, config in stages_config.items():
self.stage = stage
args = UtilsFactory.prepare_stage_args(args=args,
stage_config=config)
pprint(args)
data_params = merge_dicts(stages_data_params,
config.get("data_params", {}))
reload_loaders = data_params.get("reload_loaders", True)
if loaders is None or reload_loaders:
loaders = datasource.prepare_loaders(args=args,
stage=stage,
**data_params)
state_params = merge_dicts(stages_state_params,
config.get("state_params", {}))
callbacks_params = merge_dicts(stages_callbacks_params,
config.get("callbacks_params", {}))
config["criterion_params"] = merge_dicts(
stages_criterion_params, config.get("criterion_params", {}))
config["optimizer_params"] = merge_dicts(
stages_optimizer_params, config.get("optimizer_params", {}))
config["scheduler_params"] = merge_dicts(
stages_scheduler_params, config.get("scheduler_params", {}))
callbacks = self.prepare_callbacks(
callbacks_params=callbacks_params,
args=args,
mode="train",
stage=stage)
pprint(loaders)
pprint(callbacks)
self.run_stage_init(callbacks=callbacks)
self.criterion, self.optimizer, self.scheduler = \
UtilsFactory.create_model_stuff(
model=self.model, config=config)
start_epoch = 0 if self.state is None else self.state.epoch + 1
self.train_stage(loaders=loaders,
callbacks=callbacks,
state_params=state_params,
epochs=args.epochs,
start_epoch=start_epoch,
verbose=verbose,
logdir=args.logdir)
def on_loader_start(self, state):
if self.writer is None:
self.writer = UtilsFactory.create_tflogger(
logdir=self.logdir, name="group")
def prepare_stage_args(*, args, stage_config):
return UtilsFactory.prepare_stage_args(args=args,
stage_config=stage_config)
def create_model_stuff(*, model, config):
return UtilsFactory.create_model_stuff(model=model, config=config)