def train(epoch, dataloader):
model.train()
loss_average = AverageMeter()
error_average = ClassErrorMeter()
for i, (images, labels) in enumerate(dataloader):
n = images.size()[0]
y = model(images.cuda())
optimizer.zero_grad()
loss = criterion(y, labels.cuda())
loss_average.add(loss.item(), n)
error = error_average.add(y.cpu().data, labels)
loss.backward()
optimizer.step()
log_iter(epoch, 'train', i, loss.item(), error)
log_epoch(epoch, 'train', loss_average(), error_average())
return error_average()
def test(epoch, dataloader, val=False):
if len(dataloader) == 0:
return 1.0
mode = 'val' if val else 'test'
model.eval()
loss_average = AverageMeter()
error_average = ClassErrorMeter()
for i, (images, labels) in enumerate(dataloader):
with torch.no_grad():
n = images.size()[0]
y = model(images.cuda())
loss = criterion(y, labels.cuda())
loss_average.add(loss.item(), n)
error = error_average.add(y.cpu().data, labels)
log_iter(epoch, mode, i, loss.item(), error)
log_epoch(epoch, mode, loss_average(), error_average())
return error_average()
def main():
cl_args = parse_command_line()
config["args"].update(vars(cl_args))
# override config with command line args
if config["args"]["bases"] is not None:
assert config["args"]["bases"] in ["dct", "svd"], "Invalid bases: {}".format(config["args"]["bases"])
config["bases_to_use"] = config["args"]["bases"]
if config["args"]["input"] is not None:
assert config["args"]["input"] in ["det", "gt"], "Invalid input: {}".format(config["args"]["input"])
config["input_data"] = config["args"]["input"]
if config["args"]["nframes"] is not None:
config["n_frames"] = config["args"]["nframes"]
if config["args"]["nbases"] is not None:
config["n_bases"] = config["args"]["nbases"]
if config["args"]["gpu"] is not None:
config["gpus"] = [config["args"]["gpu"]]
# exp folder
exp_tag = "{}_F{}_k{}_{}_{}".format(datetime.datetime.now().strftime("%m%d_%H%M%S"), config["n_frames"], config["n_bases"], config["bases_to_use"], config["input_data"])
exp_path = os.path.join(config["exp_root"], exp_tag)
config["exp_tag"] = exp_tag
config["exp_path"] = exp_path
mkdir(exp_path)
# logger
logger = logging.getLogger()
coloredlogs.install(level="DEBUG", logger=logger)
fileHandler = logging.FileHandler(os.path.join(exp_path, "log.txt"))
logFormatter = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s - %(message)s")
fileHandler.setFormatter(logFormatter)
logger.addHandler(fileHandler)
config["logger"] = logger
logger.info(sys.argv)
# setup gpus
gpus = ','.join([str(x) for x in config["gpus"]])
os.environ["CUDA_VISIBLE_DEVICES"] = gpus
logger.info("Set CUDA_VISIBLE_DEVICES to {}".format(gpus))
# model
model = PoseNet(config)
model = nn.DataParallel(model)
model = model.cuda()
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=config["train"]["init_lr"])
# load bases
if config["bases_to_use"] == "svd":
fixed_bases = np.load(config["bases_path"])
assert config["n_bases"] <= fixed_bases.shape[0] and config["n_frames"] == fixed_bases.shape[1], fixed_bases.shape
fixed_bases = fixed_bases[:config["n_bases"]]
# scale svd bases to the same magnitude as dct bases
# the scaling factor here is for F=50
fixed_bases *= np.sqrt(25)
elif config["bases_to_use"] == "dct":
x = np.arange(config["n_frames"])
fixed_bases = [np.ones([config["n_frames"]]) * np.sqrt(0.5)]
for i in range(1, config["n_bases"]):
fixed_bases.append(np.cos(i * np.pi * ((x + 0.5) / config["n_frames"])))
fixed_bases = np.array(fixed_bases)
else:
assert False, config["bases_to_use"]
config["bases"] = fixed_bases
fixed_bases = torch.from_numpy(fixed_bases).float() # (K, F)
fixed_bases = fixed_bases.view(1, config["n_bases"], config["n_frames"]) # (1, K, F)
# dataset & dataloader
train_dataset = H36M_Dataset(config, "train")
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config["batch_size_per_gpu"]*len(config["gpus"]), shuffle=True, num_workers=config["num_workers"], pin_memory=True)
test_dataset = H36M_Dataset(config, "test")
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=config["batch_size_per_gpu"]*len(config["gpus"]), shuffle=False, num_workers=config["num_workers"], pin_memory=True)
tot_step = 0
for epoch in range(config["train"]["num_epochs"]):
# learning rate decay
if epoch in [60, 85]:
optimizer.param_groups[0]["lr"] = optimizer.param_groups[0]["lr"] * config["train"]["lr_decay"]
logger.info("Learning rate set to {}".format(optimizer.param_groups[0]["lr"]))
# train one epoch
model.train()
for step, batch in enumerate(train_loader):
# parse batch data
data_2d_gt = batch["data_2d_gt"] # (B, Jx2, F)
data_2d_cpn = batch["data_2d_cpn"] # (B, Jx2, F)
if config["input_data"] == "det":
data_2d = data_2d_cpn
elif config["input_data"] == "gt":
data_2d = data_2d_gt
else:
assert False, config["input_data"]
data_3d = batch["data_3d"] # (B, Jx3, F)
mean_3d = batch["mean_3d"] # (B, Jx3)
std_3d = batch["std_3d"] # (B, Jx3)
B = data_3d.shape[0]
batch_bases = fixed_bases.repeat(B, 1, 1) # (B, K, F)
data_2d = data_2d.cuda()
data_3d = data_3d.cuda()
batch_bases = batch_bases.cuda()
mean_3d = mean_3d.cuda()
std_3d = std_3d.cuda()
# forward pass
coeff = model(data_2d, batch_bases)
# compute loss
loss = model.module.build_loss_training(coeff, batch_bases, data_3d, mean_3d, std_3d)
# backward pass
optimizer.zero_grad()
loss.backward()
optimizer.step()
tot_step += 1
if "log_per_n_iterations" in config["train"] and (step + 1) % config["train"]["log_per_n_iterations"] == 0:
logger.info("TRAIN Epoch {}, step {}/{} ({}): loss = {:.6f}".format(
epoch + 1,
step + 1,
len(train_loader),
tot_step,
loss.item()))
# testing
model.eval()
logger.info("Testing on test set...")
total_loss = AverageMeter()
gts_3d = []
preds_3d = []
indices = []
with torch.no_grad():
for step, batch in enumerate(test_loader):
# parse batch data
data_2d_gt = batch["data_2d_gt"] # (B, Jx2, F)
data_2d_cpn = batch["data_2d_cpn"] # (B, Jx2, F)
if config["input_data"] == "det":
data_2d = data_2d_cpn
elif config["input_data"] == "gt":
data_2d = data_2d_gt
else:
assert False, config["input_data"]
data_2d_gt_flip = batch["data_2d_gt_flip"] # (B, Jx2, F)
data_2d_cpn_flip = batch["data_2d_cpn_flip"] # (B, Jx2, F)
if config["input_data"] == "det":
data_2d_flip = data_2d_cpn_flip
elif config["input_data"] == "gt":
data_2d_flip = data_2d_gt_flip
else:
assert False, config["input_data"]
data_3d = batch["data_3d"] # (B, Jx3, F)
data_3d_flip = batch["data_3d_flip"] # (B, Jx3, F)
mean_3d = batch["mean_3d"] # (B, Jx3)
std_3d = batch["std_3d"] # (B, Jx3)
idx = batch["idx"] # (B,)
B = data_3d.shape[0]
batch_bases = fixed_bases.repeat(B, 1, 1) # (B, K, F)
data_2d = data_2d.cuda()
data_2d_flip = data_2d_flip.cuda()
data_3d = data_3d.cuda()
data_3d_flip = data_3d_flip.cuda()
batch_bases = batch_bases.cuda()
mean_3d = mean_3d.cuda()
std_3d = std_3d.cuda()
# forward pass
coeff = model(data_2d, batch_bases)
coeff_flip = model(data_2d_flip, batch_bases)
# compute loss
loss, res = model.module.build_loss_test((coeff, coeff_flip), batch_bases, (data_3d, data_3d_flip), mean_3d, std_3d)
pred_3d, gt_3d = res
total_loss.add(loss.item())
preds_3d.append(pred_3d)
gts_3d.append(gt_3d)
indices.append(idx.data.numpy())
avg_loss = total_loss.value()
logger.info("Test loss: {}".format(avg_loss))
if epoch == config["train"]["num_epochs"] - 1:
preds_3d = np.concatenate(preds_3d, 0)
gts_3d = np.concatenate(gts_3d, 0)
indices = np.concatenate(indices, 0)
h36m_evaluate(preds_3d, gts_3d, indices, test_dataset, config)
save_ckpt(model, exp_path)
