in BigGAN_PyTorch/trainer.py [0:0]
def train(rank, world_size, config, dist_url):
print("Rank of this job is ", rank)
copy_locally = False
tmp_dir = ""
if config["ddp_train"]:
if dist_url == "env://":
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12355"
local_rank = rank
else:
rank = int(os.environ.get("SLURM_PROCID"))
local_rank = int(os.environ.get("SLURM_LOCALID"))
copy_locally = True
tmp_dir = "/scratch/slurm_tmpdir/" + str(os.environ.get("SLURM_JOB_ID"))
print("Before setting process group")
print(dist_url, rank)
dist.init_process_group(
backend="nccl", init_method=dist_url, rank=rank, world_size=world_size
)
print("After setting process group")
device = "cuda:{}".format(local_rank) # rank % 8)
print(dist_url, rank, " /Device is ", device)
else:
device = "cuda"
local_rank = "cuda"
# Update the config dict as necessary
# This is for convenience, to add settings derived from the user-specified
# configuration into the config-dict (e.g. inferring the number of classes
# and size of the images from the dataset, passing in a pytorch object
# for the activation specified as a string)'
# Seed RNG
utils.seed_rng(config["seed"] + rank)
# Setup cudnn.benchmark for free speed
torch.backends.cudnn.benchmark = True
if config["deterministic_run"]:
torch.backends.cudnn.deterministic = True
# Import the model--this line allows us to dynamically select different files.
model = __import__(config["model"])
experiment_name = (
config["experiment_name"]
if config["experiment_name"]
else utils.name_from_config(config)
)
print("Experiment name is %s" % experiment_name)
if config["ddp_train"]:
torch.cuda.set_device(device)
# Next, build the model
G = model.Generator(**{**config, "embedded_optimizers": False}).to(device)
D = model.Discriminator(**{**config, "embedded_optimizers": False}).to(device)
# If using EMA, prepare it
if config["ema"]:
print("Preparing EMA for G with decay of {}".format(config["ema_decay"]))
G_ema = model.Generator(**{**config, "skip_init": True, "no_optim": True}).to(
device
)
ema = utils.ema(G, G_ema, config["ema_decay"], config["ema_start"])
else:
G_ema, ema = None, None
print(
"Number of params in G: {} D: {}".format(
*[sum([p.data.nelement() for p in net.parameters()]) for net in [G, D]]
)
)
# Setup the optimizers
if config["D_fp16"]:
print("Using fp16 adam ")
optim_type = utils.Adam16
else:
optim_type = optim.Adam
optimizer_D = optim_type(
params=D.parameters(),
lr=config["D_lr"],
betas=(config["D_B1"], config["D_B2"]),
weight_decay=0,
eps=config["adam_eps"],
)
optimizer_G = optim_type(
params=G.parameters(),
lr=config["G_lr"],
betas=(config["G_B1"], config["G_B2"]),
weight_decay=0,
eps=config["adam_eps"],
)
# Prepare state dict, which holds things like epoch # and itr #
state_dict = {
"itr": 0,
"epoch": 0,
"save_num": 0,
"save_best_num": 0,
"best_IS": 0,
"best_FID": 999999,
"es_epoch": 0,
"config": config,
}
# FP16?
if config["G_fp16"]:
print("Casting G to float16...")
G = G.half()
if config["ema"]:
G_ema = G_ema.half()
if config["D_fp16"]:
print("Casting D to fp16...")
D = D.half()
## DDP the models
if config["ddp_train"]:
print("before G DDP ")
G = DDP(
G,
device_ids=[local_rank],
output_device=local_rank,
find_unused_parameters=True,
)
print("After G DDP ")
D = DDP(
D,
device_ids=[local_rank],
output_device=local_rank,
find_unused_parameters=True,
)
# If loading from a pre-trained model, load weights
print("Loading weights...")
if config["ddp_train"]:
dist.barrier()
map_location = device
else:
map_location = None
utils.load_weights(
G,
D,
state_dict,
config["weights_root"],
experiment_name,
config["load_weights"] if config["load_weights"] else None,
G_ema if config["ema"] else None,
map_location=map_location,
embedded_optimizers=False,
G_optim=optimizer_G,
D_optim=optimizer_D,
)
# wrapper class
GD = model.G_D(G, D, optimizer_G=optimizer_G, optimizer_D=optimizer_D)
if config["parallel"] and world_size > -1:
GD = nn.DataParallel(GD)
if config["cross_replica"]:
patch_replication_callback(GD)
# Prepare loggers for stats; metrics holds test metrics,
# lmetrics holds any desired training metrics.
if rank == 0:
test_metrics_fname = "%s/%s_log.jsonl" % (config["logs_root"], experiment_name)
train_metrics_fname = "%s/%s" % (config["logs_root"], experiment_name)
print("Inception Metrics will be saved to {}".format(test_metrics_fname))
test_log = utils.MetricsLogger(test_metrics_fname, reinitialize=False)
print("Training Metrics will be saved to {}".format(train_metrics_fname))
train_log = utils.MyLogger(
train_metrics_fname, reinitialize=False, logstyle=config["logstyle"]
)
# Write metadata
utils.write_metadata(config["logs_root"], experiment_name, config, state_dict)
else:
test_log = None
train_log = None
D_batch_size = (
config["batch_size"] * config["num_D_steps"] * config["num_D_accumulations"]
)
if config["longtail"]:
samples_per_class = np.load(
"imagenet_lt/imagenet_lt_samples_per_class.npy", allow_pickle=True
)
class_probabilities = np.load(
"imagenet_lt/imagenet_lt_class_prob.npy", allow_pickle=True
)
else:
samples_per_class, class_probabilities = None, None
train_dataset = data_utils.get_dataset_hdf5(
**{
**config,
"data_path": config["data_root"],
"batch_size": D_batch_size,
"augment": config["hflips"],
"local_rank": local_rank,
"copy_locally": copy_locally,
"tmp_dir": tmp_dir,
"ddp": config["ddp_train"],
}
)
train_loader = data_utils.get_dataloader(
**{
**config,
"dataset": train_dataset,
"batch_size": config["batch_size"],
"start_epoch": state_dict["epoch"],
"start_itr": state_dict["itr"],
"longtail_temperature": config["longtail_temperature"],
"samples_per_class": samples_per_class,
"class_probabilities": class_probabilities,
"rank": rank,
"world_size": world_size,
"shuffle": True,
"drop_last": True,
}
)
# Prepare inception metrics: FID and IS
is_moments_prefix = "I" if config["which_dataset"] == "imagenet" else "COCO"
im_filename = "%s%i_%s" % (
is_moments_prefix,
config["resolution"],
"" if not config["longtail"] else "longtail",
)
print("Using ", im_filename, "for Inception metrics.")
get_inception_metrics = inception_utils.prepare_inception_metrics(
im_filename,
samples_per_class,
config["parallel"],
config["no_fid"],
config["data_root"],
device=device,
)
G_batch_size = config["G_batch_size"]
z_, y_ = data_utils.prepare_z_y(
G_batch_size,
G.module.dim_z if config["ddp_train"] else G.dim_z,
config["n_classes"],
device=device,
fp16=config["G_fp16"],
longtail_gen=config["longtail_gen"],
custom_distrib=config["custom_distrib_gen"],
longtail_temperature=config["longtail_temperature"],
class_probabilities=class_probabilities,
)
# Balance instance sampling for ImageNet-LT
weights_sampling = None
if (
config["longtail"]
and config["use_balanced_sampler"]
and config["instance_cond"]
):
if config["which_knn_balance"] == "center_balance":
print(
"Balancing the instance features." "Using custom temperature distrib?",
config["custom_distrib_gen"],
" with temperature",
config["longtail_temperature"],
)
weights_sampling = data_utils.make_weights_for_balanced_classes(
samples_per_class,
train_loader.dataset.labels,
1000,
config["custom_distrib_gen"],
config["longtail_temperature"],
class_probabilities=class_probabilities,
)
# Balancing the NN classes (p(y))
elif config["which_knn_balance"] == "nnclass_balance":
print(
"Balancing the class distribution (classes drawn from the neighbors)."
" Using custom temperature distrib?",
config["custom_distrib_gen"],
" with temperature",
config["longtail_temperature"],
)
weights_sampling = torch.exp(
class_probabilities / config["longtail_temperature"]
) / torch.sum(
torch.exp(class_probabilities / config["longtail_temperature"])
)
# Configure conditioning sampling function to train G
sample_conditioning = functools.partial(
data_utils.sample_conditioning_values,
z_=z_,
y_=y_,
dataset=train_dataset,
batch_size=G_batch_size,
weights_sampling=weights_sampling,
ddp=config["ddp_train"],
constant_conditioning=config["constant_conditioning"],
class_cond=config["class_cond"],
instance_cond=config["instance_cond"],
nn_sampling_strategy=config["which_knn_balance"],
)
print("G batch size ", G_batch_size)
# Loaders are loaded, prepare the training function
train = train_fns.GAN_training_function(
G,
D,
GD,
ema,
state_dict,
config,
sample_conditioning,
embedded_optimizers=False,
device=device,
batch_size=G_batch_size,
)
# Prepare Sample function for use with inception metrics
sample = functools.partial(
utils.sample,
G=(G_ema if config["ema"] and config["use_ema"] else G),
sample_conditioning_func=sample_conditioning,
config=config,
class_cond=config["class_cond"],
instance_cond=config["instance_cond"],
)
print("Beginning training at epoch %d..." % state_dict["epoch"])
# Train for specified number of epochs, although we mostly track G iterations.
best_FID_run = state_dict["best_FID"]
FID = state_dict["best_FID"]
for epoch in range(state_dict["epoch"], config["num_epochs"]):
# Set epoch for distributed loader
if config["ddp_train"]:
train_loader.sampler.set_epoch(epoch)
# Initialize seeds at every epoch (useful for conditioning and
# noise sampling, as well as data order in the sampler)
if config["deterministic_run"]:
utils.seed_rng(config["seed"] + rank + state_dict["epoch"])
# Which progressbar to use? TQDM or my own?
if config["pbar"] == "mine":
pbar = utils.progress(
train_loader,
displaytype="s1k" if config["use_multiepoch_sampler"] else "eta",
)
else:
pbar = tqdm(train_loader)
s = time.time()
print("Before iteration, dataloader length", len(train_loader))
for i, batch in enumerate(pbar):
# if i> 5:
# break
in_label, in_feat = None, None
if config["instance_cond"] and config["class_cond"]:
x, in_label, in_feat, _ = batch
elif config["instance_cond"]:
x, in_feat, _ = batch
elif config["class_cond"]:
x, in_label = batch
if config["constant_conditioning"]:
in_label = torch.zeros_like(in_label)
else:
x = batch
x = x.to(device, non_blocking=True)
if in_label is not None:
in_label = in_label.to(device, non_blocking=True)
if in_feat is not None:
in_feat = in_feat.float().to(device, non_blocking=True)
# Increment the iteration counter
state_dict["itr"] += 1
# Make sure G and D are in training mode, just in case they got set to eval
# For D, which typically doesn't have BN, this shouldn't matter much.
G.train()
D.train()
if config["ema"]:
G_ema.train()
metrics = train(x, in_label, in_feat)
# print('After training step ', time.time() - s_stratified)
# s_stratified = time.time()
if rank == 0:
train_log.log(itr=int(state_dict["itr"]), **metrics)
# If using my progbar, print metrics.
if config["pbar"] == "mine" and rank == 0:
print(
", ".join(
["itr: %d" % state_dict["itr"]]
+ ["%s : %+4.3f" % (key, metrics[key]) for key in metrics]
),
end=" ",
)
# Test every specified interval
print("Iteration time ", time.time() - s)
s = time.time()
# Increment epoch counter at end of epoch
state_dict["epoch"] += 1
if not (state_dict["epoch"] % config["test_every"]):
if config["G_eval_mode"]:
print("Switching G to eval mode...")
G.eval()
D.eval()
# Compute IS and FID using training dataset as reference
test_time = time.time()
IS, FID = train_fns.test(
G,
D,
G_ema,
z_,
y_,
state_dict,
config,
sample,
get_inception_metrics,
experiment_name,
test_log,
loader=None,
embedded_optimizers=False,
G_optim=optimizer_G,
D_optim=optimizer_D,
rank=rank,
)
print("Testing took ", time.time() - test_time)
if 2 * IS < state_dict["best_IS"] and config["stop_when_diverge"]:
print("Experiment diverged!")
break
else:
print("IS is ", IS, " and 2x best is ", 2 * state_dict["best_IS"])
if not (state_dict["epoch"] % config["save_every"]) and rank == 0:
train_fns.save_weights(
G,
D,
G_ema,
state_dict,
config,
experiment_name,
embedded_optimizers=False,
G_optim=optimizer_G,
D_optim=optimizer_D,
)
if rank == 0:
if FID < best_FID_run:
best_FID_run = FID
state_dict["es_epoch"] = 0
else:
state_dict["es_epoch"] += 1
if state_dict["es_epoch"] >= config["es_patience"]:
print("reached Early stopping!")
return FID
return FID