def main(args):
flow.env.init()
flow.enable_eager_execution()
start_t = time.time()
posenet_module = PoseNet()
end_t = time.time()
print("init time : {}".format(end_t - start_t))
start_t = time.time()
pretrain_models = flow.load(args.model_path)
posenet_module.load_state_dict(pretrain_models)
end_t = time.time()
print("load params time : {}".format(end_t - start_t))
posenet_module.eval()
posenet_module.to("cuda")
start_t = time.time()
image = load_image(args.image_path)
image = flow.Tensor(image, device=flow.device("cuda"))
logits = posenet_module(image)
predictions = logits.softmax()
predictions = predictions.numpy()
end_t = time.time()
print("infer time : {}".format(end_t - start_t))
clsidx = np.argmax(predictions)
print("predict prob: %f, class name: %s" %
(np.max(predictions), clsidx_2_labels[clsidx]))
help='Output directory for video')
args = parser.parse_args()
if 'CUDA_VISIBLE_DEVICES' not in os.environ:
os.environ['CUDA_VISIBLE_DEVICES'] = args.device
settings = configparser.ConfigParser()
with open(args.config_file, 'r') as f:
settings.read_file(f)
seed = settings.getint('training', 'seed')
section = settings['hyperparameters']
dropout = section.getfloat('dropout')
# model
feature_extractor = models.resnet34(pretrained=False)
model = PoseNet(feature_extractor, droprate=dropout, pretrained=False)
model.eval()
# load weights
weights_filename = osp.expanduser(args.weights)
if osp.isfile(weights_filename):
loc_func = lambda storage, loc: storage
checkpoint = torch.load(weights_filename, map_location=loc_func)
load_state_dict(model, checkpoint['model_state_dict'])
print 'Loaded weights from {:s}'.format(weights_filename)
else:
print 'Could not load weights from {:s}'.format(weights_filename)
sys.exit(-1)
data_dir = osp.join('..', 'data', args.dataset)
stats_file = osp.join(data_dir, args.scene, 'stats.txt')
stats = np.loadtxt(stats_file)
def main(args):
flow.enable_eager_execution()
train_data_loader = OFRecordDataLoader(
ofrecord_root=args.ofrecord_path,
mode="train",
# NOTE(Liang Depeng): needs to explictly set the dataset size
dataset_size=7459,
batch_size=args.train_batch_size,
)
val_data_loader = OFRecordDataLoader(
ofrecord_root=args.ofrecord_path,
mode="val",
dataset_size=1990,
batch_size=args.val_batch_size,
)
# oneflow init
start_t = time.time()
posenet_module = PoseNet()
if args.load_checkpoint != "":
posenet_module.load_state_dict(flow.load(args.load_checkpoint))
end_t = time.time()
print("init time : {}".format(end_t - start_t))
of_cross_entropy = flow.nn.CrossEntropyLoss()
posenet_module.to("cuda")
of_cross_entropy.to("cuda")
of_sgd = flow.optim.SGD(posenet_module.parameters(),
lr=args.learning_rate,
momentum=args.mom)
of_losses = []
all_samples = len(val_data_loader) * args.val_batch_size
print_interval = 100
for epoch in range(args.epochs):
posenet_module.train()
for b in range(len(train_data_loader)):
image, label = train_data_loader.get_batch()
# oneflow train
start_t = time.time()
image = image.to("cuda")
label = label.to("cuda")
logits = posenet_module(image)
loss = of_cross_entropy(logits, label)
loss.backward()
of_sgd.step()
of_sgd.zero_grad()
end_t = time.time()
if b % print_interval == 0:
l = loss.numpy()
of_losses.append(l)
print(
"epoch {} train iter {} oneflow loss {}, train time : {}".
format(epoch, b, l, end_t - start_t))
print("epoch %d train done, start validation" % epoch)
posenet_module.eval()
correct_of = 0.0
for b in range(len(val_data_loader)):
image, label = val_data_loader.get_batch()
start_t = time.time()
image = image.to("cuda")
with flow.no_grad():
logits = posenet_module(image)
predictions = logits.softmax()
of_predictions = predictions.numpy()
clsidxs = np.argmax(of_predictions, axis=1)
label_nd = label.numpy()
for i in range(args.val_batch_size):
if clsidxs[i] == label_nd[i]:
correct_of += 1
end_t = time.time()
print("epoch %d, oneflow top1 val acc: %f" %
(epoch, correct_of / all_samples))
flow.save(
posenet_module.state_dict(),
os.path.join(
args.save_checkpoint_path,
"epoch_%d_val_acc_%f" % (epoch, correct_of / all_samples),
),
)
writer = open("of_losses.txt", "w")
for o in of_losses:
writer.write("%f\n" % o)
writer.close()