def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
print('Creating model...')
self.model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
self.model, _, start_epoch = load_model(self.model, opt.load_model,
opt)
self.model = self.model.to(opt.device)
self.model.eval()
self.start_epoch = start_epoch
self.opt = opt
self.trained_dataset = get_dataset(opt.dataset)
self.mean = np.array(self.trained_dataset.mean,
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array(self.trained_dataset.std,
dtype=np.float32).reshape(1, 1, 3)
self.pause = not opt.no_pause
self.rest_focal_length = self.trained_dataset.rest_focal_length \
if self.opt.test_focal_length < 0 else self.opt.test_focal_length
self.flip_idx = self.trained_dataset.flip_idx
self.cnt = 0
self.pre_images = None
self.pre_image_ori = None
self.age_images = []
self.tracker = SchTracker(opt) if opt.sch_track else Tracker(opt)
self.debugger = Debugger(opt=opt, dataset=self.trained_dataset)
self.decoder = GenericDecode(K=opt.K, opt=opt)
def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
print('Creating model...')
self.model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
self.model = load_model(self.model, opt.load_model, opt)
self.model = self.model.to(opt.device)
self.model.eval()
# inp = (torch.ones([1, 3, 320, 320]).cuda(),
# torch.ones([1, 3, 320, 320]).cuda(),
# torch.ones([1, 1, 320, 320]).cuda())
# pytorch_to_caffe.trans_net(self.model, inp, 'res18')
# pytorch_to_caffe.save_prototxt('{}.prototxt'.format('res18'))
# pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format('res18'))
self.opt = opt
self.trained_dataset = get_dataset(opt.dataset)
self.mean = np.array(self.trained_dataset.mean,
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array(self.trained_dataset.std,
dtype=np.float32).reshape(1, 1, 3)
self.pause = not opt.no_pause
self.rest_focal_length = self.trained_dataset.rest_focal_length \
if self.opt.test_focal_length < 0 else self.opt.test_focal_length
self.flip_idx = self.trained_dataset.flip_idx
self.cnt = 0
self.pre_images = None
self.pre_image_ori = None
self.tracker = Tracker(opt)
self.debugger = Debugger(opt=opt, dataset=self.trained_dataset)
def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
print('Creating model...')
self.model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
self.model = load_model(self.model, opt.load_model, opt)
self.model = self.model.to(opt.device)
self.model.eval()
self.opt = opt
self.trained_dataset = get_dataset(opt.dataset)
self.mean = np.array(self.trained_dataset.mean,
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array(self.trained_dataset.std,
dtype=np.float32).reshape(1, 1, 3)
self.pause = not opt.no_pause
self.rest_focal_length = self.trained_dataset.rest_focal_length \
if self.opt.test_focal_length < 0 else self.opt.test_focal_length
self.flip_idx = self.trained_dataset.flip_idx
self.cnt = 0
self.pre_images = None
self.pre_image_ori = None
self.inference_feats = deque(maxlen=self.opt.clip_len - 1)
self.inference_preim = deque(maxlen=self.opt.clip_len - 1)
self.inference_prehm = deque(maxlen=self.opt.clip_len - 1)
self.tracker = Tracker(opt)
self.debugger = Debugger(opt=opt, dataset=self.trained_dataset)
def convert_onnx(opt):
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
opt.model_output_list = True
if opt.gpus[0] >= 0:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
Dataset = dataset_factory[opt.test_dataset]
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
print(opt)
model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
if opt.load_model != '':
model = load_model(model, opt.load_model, opt)
model = model.to(opt.device)
model.eval()
dummy_input1 = torch.randn(1, 3, opt.input_h, opt.input_w).to(opt.device)
if opt.tracking:
dummy_input2 = torch.randn(1, 3, opt.input_h,
opt.input_w).to(opt.device)
if opt.pre_hm:
dummy_input3 = torch.randn(1, 1, opt.input_h,
opt.input_w).to(opt.device)
torch.onnx.export(model,
(dummy_input1, dummy_input2, dummy_input3),
"../models/{}.onnx".format(opt.exp_id))
else:
torch.onnx.export(model, (dummy_input1, dummy_input2),
"../models/{}.onnx".format(opt.exp_id))
else:
torch.onnx.export(model, (dummy_input1, ),
"../models/{}.onnx".format(opt.exp_id))
def run(path, batch_size=4, debug=0):
pre_model_load = time.time()
model = create_model()
model = load_model(model, 'checkpoints/coco_tracking.pth')
model.to(torch.device('cuda'))
model.eval()
del model
model_loading_time = time.time() - pre_model_load
multi_vid_list = get_sorted_list(path=path)
q_in_model = Queue(10)
q_out_model = Queue(10)
q_times = Queue(batch_size + 10)
reader_thread = Thread(target=reader_thread_fn,
args=(q_in_model, q_times, batch_size, path,
multi_vid_list))
model_thread = Thread(target=model_thread_fn,
args=(q_in_model, q_out_model))
tracker_thread = Thread(target=tracker_thread_fn,
args=(q_out_model, q_times, multi_vid_list,
model_loading_time, batch_size, debug))
reader_thread.start()
model_thread.start()
tracker_thread.start()
reader_thread.join()
model_thread.join()
tracker_thread.join()
print("Total time: {}".format(time.time() - pre_model_load),
file=sys.stderr)
def run_expreminent(look_forward, hidden_size, batch_size, epochs, dropout,
dataset):
x, y = get_sentiments_prices(dataset['twitter_sentiments'],
dataset["reddit_sentiments"],
dataset["coin_price"], look_forward)
for i in range(x.shape[1]):
x[:, i] = normalize_array(x[:, i])
# split into train and test sets
train_x, test_x = split(x)
train_y, test_y = split(y)
train_x = np.reshape(train_x,
(train_x.shape[0], look_forward, train_x.shape[1]))
test_x = np.reshape(test_x,
(test_x.shape[0], look_forward, test_x.shape[1]))
model = create_model(hidden_size=hidden_size,
look_forward=look_forward,
dropout=dropout)
model = train(model,
train_x,
train_y,
batch_size=batch_size,
epochs=epochs)
y_pred = test(model, test_x)
score = evaluate(test_y, y_pred)
print('Test Score: %.2f RMSE' % score)
return score
def train(args):
env = gym.make(args.env)
num_actions = env.action_space.n
model = create_model(
(args.cropped_size, args.cropped_size, args.stack_frames), num_actions)
processor = ProcessorProxy((args.cropped_size, args.cropped_size))
memory = ExperienceMemory(num_frames=args.stack_frames,
max_size=args.memsize)
policy = PolicyProxy(num_actions, config.epsilon_greedy)
dqn = args.dqn
if dqn is None:
dqn = config.nature_dqn
agent = AgentProxy(agent_type=dqn,
model=model,
processor=processor,
memory=memory,
policy=policy,
gamma=args.gamma,
batch_size=args.batch_size,
target_update_freq=args.target_update_freq,
num_burn_in=args.num_burn_in,
train_freq=args.train_freq,
num_iterations=args.num_iterations,
save_path=args.output,
max_episode_length=args.max_episode_length)
agent.load_weights(args.model_path)
agent.compile(optimizer=Adam(lr=args.learning_rate), loss_func=dqn_loss)
sys.stdout.flush()
agent.fit(env)
def train(num_epochs, save_dir, dataset, restore):
data = dataset_utils.load_dataset(dataset)
split = 0.1
train_split, validation_split, batch_size = dataset_utils.create_generator_and_validation(data, split)
model, optimizer, loss_function = models.create_model()
max_to_keep = 25
checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
manager = tf.train.CheckpointManager(checkpoint, directory='checkpoints', max_to_keep=max_to_keep)
best_checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
best_manager = tf.train.CheckpointManager(best_checkpoint, directory='checkpoints/best', max_to_keep=1)
log_dir = f'logs/{datetime.now().strftime("%d-%b-%Y_%H.%M.%S")}'
callbacks=[
tf.keras.callbacks.EarlyStopping(patience=num_epochs*0.1),
tf.keras.callbacks.TensorBoard(log_dir=log_dir, write_graph=False),
CheckpointCallback(manager, best_manager)
]
if restore:
print('Restoring from latest checkpoint.')
manager.restore_or_initialize()
model.fit(x=train_split, epochs=num_epochs, batch_size=batch_size, validation_data=validation_split, callbacks=callbacks)
print(f'TensorBoard logs can be found under \'{log_dir}\'.')
best_manager.restore_or_initialize()
print(f'Saving model as \'{save_dir}\'.')
model.save(save_dir)
def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
print('Creating model...')
self.model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
self.model = load_model(self.model, opt.load_model, opt)
self.model = self.model.to(opt.device)
self.model.eval()
self.opt = opt
self.trained_dataset = get_dataset(opt.dataset)
self.mean = np.array(self.trained_dataset.mean,
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array(self.trained_dataset.std,
dtype=np.float32).reshape(1, 1, 3)
self.pause = not opt.no_pause
self.rest_focal_length = self.trained_dataset.rest_focal_length \
if self.opt.test_focal_length < 0 else self.opt.test_focal_length
self.flip_idx = self.trained_dataset.flip_idx
self.cnt = 0
self.pre_images = None
self.pre_image_ori = None
self.tracker = Tracker(opt)
self.debugger = Debugger(opt=opt, dataset=self.trained_dataset)
self.motion = opt.motion
if self.motion == 'transformer':
import sys
M3_PATH = '/u/jozhang/code/motion3d/'
sys.path.insert(0, M3_PATH)
from models.transformer import DPTransformer
# motion = DPTransformer(2, 64, {'depth': 3, 'heads': 8, 'dim_head': 8, 'mlp_dim': 64, 'dropout': 0.})
# trans_path = '/scratch/cluster/jozhang/logs/hydra/2021-01-30/15-36-54/models/ckpt-latest.dat'
ckpt = torch.load(opt.transformer_load_path)
self.transformer = ckpt['model'].cuda()
print(
f'Using transformer motion loaded from {opt.transformer_load_path}'
)
elif self.motion == 'zero':
print(f'Using no motion model')
elif self.motion == 'cttrack':
print(f'Using cttrack motion model')
else:
assert False, f'Do not recognize such motion model {self.motion}'
self.negate_motion = opt.negate_motion
if self.negate_motion:
logging.warning('Motion is being negated! Are you sure?')
self.all_pre_images = []
def start_training(cuda, epochs, general_seed, tensorflow_seed, batch_size,
buffer_size, learning_rate):
# Disable GPU support if no GPUs are supposed to be used
if not cuda:
tf.config.set_visible_devices([], 'GPU')
with mlflow.start_run():
# Enable the logging of all parameters, metrics and models to mlflow and Tensorboard
mlflow.tensorflow.autolog()
# Fix all random seeds and Tensorflow specific reproducibility settings
set_general_random_seeds(general_seed)
set_tensorflow_random_seeds(tensorflow_seed)
# Use Mirrored Strategy for multi GPU support
strategy = tf.distribute.MirroredStrategy()
click.echo(
click.style(f'Number of devices: {strategy.num_replicas_in_sync}',
fg='blue'))
# Fetch and prepare dataset
train_dataset, eval_dataset = load_train_test_data(
strategy, batch_size, buffer_size, tensorflow_seed)
with strategy.scope():
# Define model and compile model
model = create_model(input_shape=(28, 28, 1))
model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True),
optimizer=tf.keras.optimizers.Adam(
learning_rate=learning_rate),
metrics=['accuracy'])
# Train and evaluate the trained model
runtime = time.time()
train(model, epochs, train_dataset)
eval_loss, eval_acc = test(model, eval_dataset)
click.echo(f'Test loss: {eval_loss}, Test Accuracy: {eval_acc}')
device = 'GPU' if cuda else 'CPU'
click.echo(
click.style(
f'{device} Run Time: {str(time.time() - runtime)} seconds',
fg='green'))
# Log hardware and software
log_sys_intel_conda_env()
click.echo(
click.style(
f'\nLaunch TensorBoard with:\ntensorboard --logdir={os.path.join(mlflow.get_artifact_uri(), "tensorboard_logs", "train")}',
fg='blue'))
def main():
model = create_model(trainable=TRAINABLE)
# if TRAINABLE:
# model.load_weights(WEIGHTS)
train_datagen = DataGenerator(file_path=cfg.TRAIN.DATA_PATH,
config_path=cfg.TRAIN.ANNOTATION_PATH)
val_generator = DataGenerator(file_path=cfg.TEST.DATA_PATH,
config_path=cfg.TEST.ANNOTATION_PATH,
debug=False)
validation_datagen = Validation(generator=val_generator)
learning_rate = cfg.TRAIN.LEARNING_RATE
if TRAINABLE:
learning_rate /= 10
optimizer = tf.keras.optimizers.SGD(lr=learning_rate,
decay=cfg.TRAIN.LR_DECAY,
momentum=0.9,
nesterov=False)
model.compile(loss=detect_loss(), optimizer=optimizer, metrics=[])
checkpoint = tf.keras.callbacks.ModelCheckpoint("model-{val_iou:.2f}.h5",
monitor="val_iou",
verbose=1,
save_best_only=True,
save_weights_only=True,
mode="max")
stop = tf.keras.callbacks.EarlyStopping(monitor="val_iou",
patience=cfg.TRAIN.PATIENCE,
mode="max")
reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor="val_iou",
factor=0.6,
patience=5,
min_lr=1e-6,
verbose=1,
mode="max")
# Define the Keras TensorBoard callback.
logdir = "logs/fit/" + datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)
model.fit_generator(generator=train_datagen,
epochs=cfg.TRAIN.EPOCHS,
callbacks=[
tensorboard_callback, validation_datagen,
checkpoint, reduce_lr, stop
],
shuffle=True,
verbose=1)
def start_training(cuda, epochs, general_seed, pytorch_seed, log_interval,
training_batch_size, test_batch_size, learning_rate):
# Set GPU settings
use_cuda = (True if cuda == 'True' else False) and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
if use_cuda and torch.cuda.device_count() > 0:
click.echo(click.style(f'Using {torch.cuda.device_count()} GPUs!', fg='blue'))
# Set all random seeds and possibly turn of GPU non determinism
set_general_random_seeds(general_seed)
set_pytorch_random_seeds(pytorch_seed, use_cuda=use_cuda)
# Load training and testing data
train_loader, test_loader = load_train_test_data(training_batch_size, test_batch_size)
# Define model, device and optimizer
if torch.cuda.device_count() > 1:
model = create_parallel_model()
else:
model = create_model()
model.to(device)
optimizer = optim.Adam(model.parameters())
optimizer.step()
with mlflow.start_run():
# Create a SummaryWriter to write TensorBoard events locally
events_output_dir = tempfile.mkdtemp()
writer = SummaryWriter(events_output_dir)
click.echo(click.style(f'Writing TensorBoard events locally to {events_output_dir}\n', fg='blue'))
# Start training
runtime = time.time()
for epoch in range(1, epochs + 1):
train(use_cuda, model, epoch, optimizer, log_interval, train_loader, writer)
test(use_cuda, model, epoch, test_loader, writer)
device = 'GPU' if use_cuda else 'CPU'
click.echo(click.style(f'{device} Run Time: {str(time.time() - runtime)} seconds', fg='green'))
# Closing writer to allow for the model to be logged
writer.close()
# Log the model to mlflow
click.echo(click.style('Logging model to mlflow...', fg='blue'))
mlflow.pytorch.log_model(model, 'models')
# Log hardware and software
log_sys_intel_conda_env()
# Upload the TensorBoard event logs as a run artifact
click.echo(click.style('Uploading TensorBoard events as a run artifact...', fg='blue'))
mlflow.log_artifacts(events_output_dir, artifact_path='events')
click.echo(click.style(f'\nLaunch TensorBoard with:\ntensorboard --logdir={os.path.join(mlflow.get_artifact_uri(), "events")}', fg='blue'))
def run(path, batch_size=4, debug=0):
pre_model_load = time.time()
model = create_model()
model = load_model(model, 'checkpoints/coco_tracking.pth')
model.to(torch.device('cuda'))
model.eval()
model_loading_time = time.time() - pre_model_load
multi_vid_list = get_sorted_list(path=path)
vid_managers = [VideoManager(path, vid_list, model_loading_time) for vid_list in multi_vid_list[:batch_size]]
next_video_id = len(vid_managers)
processed_frames = 0
done = False
while len(vid_managers) > 0:
imgs = [manager.get_img() for manager in vid_managers]
cur_imgs = torch.cat([x[0] for x in imgs], dim=0)
prev_imgs = torch.cat([x[1] for x in imgs], dim=0)
with torch.no_grad():
with torch.cuda.amp.autocast(enabled=True):
out = model(cur_imgs, prev_imgs, None)[-1]
out = sigmoid_output(out)
dets = generic_decode(out)
for k in dets:
dets[k] = dets[k].detach().cpu().numpy()
processed_frames += len(vid_managers)
for i, manager in reversed(list(enumerate(vid_managers))):
single_dets = {k: v[np.newaxis, i] for k, v in dets.items()}
manager.process_output(single_dets)
if manager.is_done():
manager.finalize()
if next_video_id < len(multi_vid_list):
vid_managers[i] = VideoManager(path, multi_vid_list[next_video_id], model_loading_time)
next_video_id += 1
else:
del vid_managers[i]
if debug:
frame_time = time.time() - pre_model_load
FPS = processed_frames / frame_time
print("At frame {} FPS {}".format(processed_frames, FPS), file=sys.stderr)
def main(opt):
print('Creating model...')
torch.backends.cudnn.benchmark = not opt.not_cuda_benchmark and not opt.test
Dataset = get_dataset(opt.dataset)
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
print(opt)
if not opt.not_set_cuda_env:
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
opt.device = torch.device('cuda' if opt.gpus[0] >= 0 else 'cpu')
model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
optimizer = get_optimizer(opt, model)
if opt.load_model != '':
model, optimizer, start_epoch = load_model(
model, opt.load_model, opt, optimizer)
def model_thread_fn(q_in, q_out):
model = create_model()
model = load_model(model, 'checkpoints/coco_tracking.pth')
model.to(torch.device('cuda'))
model.eval()
while True:
cur_imgs, prev_imgs = q_in.get(timeout=30)
if cur_imgs is None:
break
with torch.no_grad():
with torch.cuda.amp.autocast(enabled=True):
out = model(cur_imgs, prev_imgs, None)[-1]
out = sigmoid_output(out)
dets = generic_decode(out)
q_out.put(dets)
def evaluate(args):
env = gym.make(args.env)
num_actions = env.action_space.n
model = create_model(
(args.cropped_size, args.cropped_size, args.stack_frames), num_actions)
rewards = []
lens = []
tries = 0
while True:
env = gym.make(args.env)
env = wrappers.Monitor(env, 'videos', force=True)
processor = ProcessorProxy((args.cropped_size, args.cropped_size))
memory = ExperienceMemory(num_frames=args.stack_frames,
max_size=args.memsize)
policy = PolicyProxy(num_actions, config.greedy)
dqn = args.dqn
if dqn is None:
dqn = config.nature_dqn
agent = AgentProxy(agent_type=dqn,
model=model,
processor=processor,
memory=memory,
policy=policy,
gamma=args.gamma,
batch_size=args.batch_size,
target_update_freq=args.target_update_freq,
num_burn_in=args.num_burn_in,
train_freq=args.train_freq,
num_iterations=args.num_iterations,
save_path=args.output,
max_episode_length=args.max_episode_length)
agent.load_weights(args.model_path)
cumulative_reward, std, average_episode_length = agent.evaluate(
env=env, num_episodes=1)
tries += 1
# Sometime the model is not very stable.
if tries > 100 or cumulative_reward > 350:
break
print(
'average reward = %f, cumulative_reward=%f, std = %f, average_epis_length = %d'
% (cumulative_reward, cumulative_reward, std,
average_episode_length))
rewards.append(cumulative_reward)
lens.append(average_episode_length)
def setup_model(self, model):
global status_data
global my_model
status_data['text'] = "Loading data..."
(x_train, y_train), (x_test,
y_test) = tf.keras.datasets.mnist.load_data()
print(type(x_train))
print(x_train.shape)
print(type(x_train[0]))
print(x_train[0].shape)
x_train_normalized = x_train / 255
x_test_normalized = x_test / 255
status_data['text'] = "Training model..."
# The following variables are the hyperparameters.
learning_rate = 0.003
epochs = 50
batch_size = 4000
validation_split = 0.2
# Establish the model's topography.
my_model = model.create_model(learning_rate)
# Train the model on the normalized training set.
epochs, hist = model.train_model(my_model, x_train_normalized, y_train,
epochs, batch_size, validation_split)
# Plot a graph of the metric vs. epochs.
list_of_metrics_to_plot = ['accuracy']
model.plot_curve(epochs, hist, list_of_metrics_to_plot)
# Evaluate against the test set.
print("\n Evaluate the new model against the test set:")
my_model.evaluate(x=x_test_normalized, y=y_test, batch_size=batch_size)
status_data['text'] = "Model trained."
def main(_argv):
model = create_model()
model.load_weights(FLAGS.weights)
proc_image = tf.keras.preprocessing.image.load_img(
FLAGS.image, target_size=(cfg.NN.INPUT_SIZE, cfg.NN.INPUT_SIZE))
proc_image = tf.keras.preprocessing.image.img_to_array(proc_image)
proc_image = np.expand_dims(proc_image, axis=0)
proc_image - tf.keras.applications.mobilenet_v2.preprocess_input(
proc_image)
original_image = cv2.imread(FLAGS.image)
t1 = time.time()
pred = np.squeeze(model.predict(proc_image))
t2 = time.time()
processing_time = t2 - t1
height, width, y_f, x_f, score = [
a.flatten() for a in np.split(pred, pred.shape[-1], axis=-1)
]
coords = np.arange(pred.shape[0] * pred.shape[1])
y = (y_f + coords // pred.shape[0]) / pred.shape[0]
x = (x_f + coords % pred.shape[1]) / pred.shape[1]
boxes = np.stack([y, x, height, width, score], axis=-1)
boxes = boxes[np.where(boxes[..., -1] >= SCORE_THRESHOLD)]
# does not work with TF GPU, uncomment only when using CPU
# selected_indices = tf.image.non_max_suppression(boxes[..., :-1], boxes[..., -1], MAX_OUTPUT_SIZE, cfg.NN.IOU_LOSS_THRESH)
# selected_indices = tf.Session().run(selected_indices)
# print(len(boxes))
original_image = draw_outputs(original_image, boxes)
original_image = cv2.putText(original_image,
"Time: {:.2f}".format(processing_time),
(0, 30), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1,
(0, 0, 255), 2)
cv2.imwrite(FLAGS.output, original_image)
def setup(opt, checkpoint):
"""Create new model or reload from checkpoint model"""
# Resume model if checkpoint is provided, create a new model otherwise.
if checkpoint is not None:
model_path = os.path.join(opt.resume, checkpoint['model_file'])
utils.check_file(model_path)
print("".ljust(4) + "=> Resuming model from %s" % model_path)
model = torch.load(model_path)
else:
print("".ljust(4) + "=> Creating new model")
model = create_model(opt)
# Load optim_file if checkpoint is provided, return None otherwise.
if checkpoint is not None:
optim_path = os.path.join(opt.resume, checkpoint['optim_file'])
utils.check_file(optim_path)
print("".ljust(4) + "=> Resuming optim_state from %s" % optim_path)
optim_state = torch.load(optim_path)
else:
optim_state = None
return model, optim_state
def train():
device = torch.device('cuda' if cfg.GPU[0] >= 0 else 'cpu')
start_epoch = 1
if start_epoch == 1:
train_log = open(os.path.join(cfg.LOG_DIR, "train_log.csv"), 'w')
train_log_title = "epoch,total_loss,classify_loss,angle_loss,iou_loss\n"
train_log.write(train_log_title)
train_log.flush()
else:
train_log = open(os.path.join(cfg.LOG_DIR, "train_log.csv"), 'a')
print('Creating model...')
model = create_model()
if start_epoch != 1:
model = load_model(
model, 'logs/weights/model_epoch_{}.pth'.format(start_epoch - 1))
optimizer = torch.optim.Adam(model.parameters(), cfg.LR)
trainer = Trainer(model, optimizer)
trainer.set_device(device)
print('Setting up data...')
train_loader = DataLoader(LatexDataset(),
batch_size=cfg.BATCH_SIZE,
shuffle=True,
num_workers=cfg.NUM_WORKERS,
pin_memory=True,
drop_last=True)
print('Starting training...')
epoch = start_epoch
for epoch in range(start_epoch, start_epoch + cfg.EPOCHS):
trainer.train(epoch, train_loader, train_log)
if epoch % 5 == 0:
save_model('logs/weights/model_epoch_{}.pth'.format(epoch), epoch,
model)
save_model(os.path.join(cfg.WEIGHTS_DIR, 'model_last.pth'), epoch, model)
def __init__(self, opt):
if opt.gpus[0] >= 0:
opt.device = torch.device("cuda")
else:
opt.device = torch.device("cpu")
print("Creating model...")
self.model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
self.model = load_model(self.model, opt.load_model, opt)
self.model = self.model.to(opt.device)
self.model.eval()
self.opt = opt
self.trained_dataset = get_dataset(opt.dataset)
self.mean = np.array(self.trained_dataset.mean,
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array(self.trained_dataset.std,
dtype=np.float32).reshape(1, 1, 3)
# self.pause = not opt.no_pause
self.rest_focal_length = (self.trained_dataset.rest_focal_length
if self.opt.test_focal_length < 0 else
self.opt.test_focal_length)
self.flip_idx = self.trained_dataset.flip_idx
self.cnt = 0
self.pre_images = None
self.pre_image_ori = None
self.dataset = opt.dataset
if self.dataset == "nuscenes":
self.tracker = {}
for class_name in NUSCENES_TRACKING_NAMES:
self.tracker[class_name] = Tracker(opt, self.model)
else:
self.tracker = Tracker(opt, self.model)
self.debugger = Debugger(opt=opt, dataset=self.trained_dataset)
self.img_height = 100
self.img_width = 100
def main(opt):
torch.manual_seed(opt.seed)
torch.backends.cudnn.benchmark = not opt.not_cuda_benchmark and not opt.test
Dataset = get_dataset(opt.dataset)
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
print(opt)
if not opt.not_set_cuda_env:
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
opt.device = torch.device('cuda' if opt.gpus[0] >= 0 else 'cpu')
logger = Logger(opt)
print('Creating model...')
model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
optimizer = get_optimizer(opt, model)
start_epoch = 0
if opt.load_model != '':
model, optimizer, start_epoch = load_model(model, opt.load_model, opt,
optimizer)
############################################3333
#freezing backbone and one head
for param in model.parameters():
# print(param)
param.requires_grad = False
req_grad = ["model.hm_bdd", "model.wh_bdd", "model.reg_bdd"]
# for hd in model.reg_tl:
for custom_head in (req_grad):
for hd in eval(custom_head):
# print(hd.parameters())
for wt in hd.parameters():
# print(wt)
wt.requires_grad = True
######################################################
trainer = Trainer(opt, model, optimizer)
trainer.set_device(opt.gpus, opt.chunk_sizes, opt.device)
if opt.val_intervals < opt.num_epochs or opt.test:
print('Setting up validation data...')
val_loader = torch.utils.data.DataLoader(Dataset(opt, 'val'),
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True)
if opt.test:
_, preds = trainer.val(0, val_loader)
val_loader.dataset.run_eval(preds, opt.save_dir)
return
print('Setting up train data...')
train_loader = torch.utils.data.DataLoader(Dataset(opt, 'train'),
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True,
drop_last=True)
print('Starting training...')
for epoch in range(start_epoch + 1, opt.num_epochs + 1):
mark = epoch if opt.save_all else 'last'
log_dict_train, _ = trainer.train(epoch, train_loader)
logger.write('epoch: {} |'.format(epoch))
for k, v in log_dict_train.items():
logger.scalar_summary('train_{}'.format(k), v, epoch)
logger.write('{} {:8f} | '.format(k, v))
if opt.val_intervals > 0 and epoch % opt.val_intervals == 0:
save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(mark)),
epoch, model, optimizer)
with torch.no_grad():
log_dict_val, preds = trainer.val(epoch, val_loader)
if opt.eval_val:
val_loader.dataset.run_eval(preds, opt.save_dir)
for k, v in log_dict_val.items():
logger.scalar_summary('val_{}'.format(k), v, epoch)
logger.write('{} {:8f} | '.format(k, v))
else:
save_model(os.path.join(opt.save_dir, 'model_last.pth'), epoch,
model, optimizer)
logger.write('\n')
# if epoch in opt.save_point:
if epoch % opt.save_point[0] == 0:
save_model(
os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)),
epoch, model, optimizer)
if epoch in opt.lr_step:
lr = opt.lr * (0.1**(opt.lr_step.index(epoch) + 1))
print('Drop LR to', lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
logger.close()
def main(opt):
torch.manual_seed(opt.seed)
torch.backends.cudnn.benchmark = not opt.not_cuda_benchmark and not opt.test
Dataset = get_dataset(opt.dataset)
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
print(opt)
if not opt.not_set_cuda_env:
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus_str
opt.device = torch.device("cuda" if opt.gpus[0] >= 0 else "cpu")
logger = Logger(opt)
print("Creating model...")
model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
optimizer = get_optimizer(opt, model)
start_epoch = 0
if opt.load_model != "":
model, optimizer, start_epoch = load_model(
model, opt.load_model, opt, optimizer
)
for i, param in enumerate(model.parameters()):
param.requires_grad = True
trainer = Trainer(opt, model, optimizer)
trainer.set_device(opt.gpus, opt.chunk_sizes, opt.device)
if opt.val_intervals < opt.num_epochs or opt.test:
print("Setting up validation data...")
val_loader = torch.utils.data.DataLoader(
Dataset(opt, "val"),
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True,
)
if opt.test:
_, preds = trainer.val(0, val_loader)
val_loader.dataset.run_eval(preds, opt.save_dir)
return
print("Setting up train data...")
train_loader = torch.utils.data.DataLoader(
Dataset(opt, "train"),
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True,
drop_last=True,
)
print("Starting training...")
for epoch in range(start_epoch + 1, opt.num_epochs + 1):
save_model(
os.path.join(opt.save_dir, "model_{}.pth".format(epoch)),
epoch,
model,
optimizer,
)
mark = epoch if opt.save_all else "last"
log_dict_train, _ = trainer.train(epoch, train_loader)
logger.write("epoch: {} |".format(epoch))
for k, v in log_dict_train.items():
logger.scalar_summary("train_{}".format(k), v, epoch)
logger.write("{} {:8f} | ".format(k, v))
if opt.val_intervals > 0 and epoch % opt.val_intervals == 0:
save_model(
os.path.join(opt.save_dir, "model_{}.pth".format(mark)),
epoch,
model,
optimizer,
)
with torch.no_grad():
log_dict_val, preds = trainer.val(epoch, val_loader)
if opt.eval_val:
val_loader.dataset.run_eval(preds, opt.save_dir)
for k, v in log_dict_val.items():
logger.scalar_summary("val_{}".format(k), v, epoch)
logger.write("{} {:8f} | ".format(k, v))
else:
save_model(
os.path.join(opt.save_dir, "model_last.pth"), epoch, model, optimizer
)
logger.write("\n")
# if epoch in opt.save_point:
save_model(
os.path.join(opt.save_dir, "model_{}.pth".format(epoch)),
epoch,
model,
optimizer,
)
if epoch in opt.lr_step:
lr = opt.lr * (0.1 ** (opt.lr_step.index(epoch) + 1))
print("Drop LR to", lr)
for param_group in optimizer.param_groups:
param_group["lr"] = lr
logger.close()
def run_single_video_serial(path, debug=0, full_precision=False):
init_time = time.time()
if debug >= 1:
print("Starting for video: {}".format(path), file=sys.stderr)
video_id, camera_id, max_frames, width, height = get_video_params(path)
cap = cv2.VideoCapture(path)
model = create_model()
model = load_model(model, 'checkpoints/coco_tracking.pth')
model.to(torch.device('cuda'))
model.eval()
tracker = Tracker(init_time,
video_id,
max_frames,
camera_id,
width,
height,
debug=debug)
preprocess_function = get_img_transform(height, width, new_size=512)
postprocess_trans = get_postprocess_trans(height, width)
region_mask = get_region_mask(camera_id, height, width)
region_mask = np.where(region_mask, 255, 0).astype(np.uint8)
if debug > 2:
cv2.imwrite("mask.png", region_mask)
pre_img = None
for i in range(max_frames):
ret, frame = cap.read()
if debug >= 2:
cv2.imshow("Frame", frame)
cv2.waitKey(1)
tracker.frame = np.copy(frame)
frame = cv2.bitwise_and(frame, frame, mask=region_mask)
img = preprocess_function(frame)
img = torch.from_numpy(img).to(torch.device('cuda'))
if pre_img is None:
pre_img = img
with torch.no_grad():
with torch.cuda.amp.autocast(enabled=not full_precision):
out = model(img, pre_img, None)[-1]
out = sigmoid_output(out)
dets = generic_decode(out)
pre_img = img
for k in dets:
dets[k] = dets[k].detach().cpu().numpy()
dets = post_process(dets, postprocess_trans)[0]
tracker.step(dets)
if debug >= 1 and i % 100 == 99:
frame_time = time.time() - init_time
FPS = (i + 1) / frame_time
print("At frame {} FPS {}".format(i + 1, FPS), file=sys.stderr)
tracker.finalize()
if debug >= 1:
print("Finished video: {}".format(path), file=sys.stderr)
tf.logging.set_verbosity(tf.logging.ERROR)
app = Flask(__name__, static_url_path="/client")
IMAGE_FOLDER = "evaluate/"
DATASET_FOLDER = "dataset/"
TURING_FILES_FOLDER = "turingFiles/"
LABELS_FILE = "labels.txt"
FILE_EXTENSION = ".png"
cors = CORS(app)
app.config["CORS_HEADERS"] = "Content-Type"
print("Loading the model...")
with open("model/params.json", "r") as file:
params = json.load(file)
model, decoders = create_model(params, gpu=False)
model.load_weights("model/weights.h5")
decoder = decoders[0]
# A tensor where all values are the same, is required by ctc loss
ctc_input_length = (
params["img_w"] //
(params["pool_size"]**params["num_convs"])) - params["ctc_cut"]
# Create necessary folders if its the first time
for folder in ["dataset", "turingFiles", "evaluate"]:
if not os.path.exists(folder):
print(f"Creating folder {folder}.")
os.mkdir(folder)
print("Ready!")
def evaluate(config, mode):
print(toGreen('Loading checkpoint manager...'))
print(config.LOG_DIR.ckpt, mode)
ckpt_manager = CKPT_Manager(config.LOG_DIR.ckpt, mode, 10)
date = datetime.datetime.now().strftime('%Y.%m.%d.(%H%M)')
print(toYellow('======== EVALUATION START ========='))
##################################################
## DEFINE MODEL
print(toGreen('Initializing model'))
model = create_model(config)
model.eval()
model.print()
inputs = {
'inp': None,
'ref': None,
'inp_hist': None,
'ref_hist': None,
'seg_inp': None,
'seg_ref': None
}
inputs = collections.OrderedDict(sorted(inputs.items(),
key=lambda t: t[0]))
## INITIALIZING VARIABLE
print(toGreen('Initializing variables'))
result, ckpt_name = ckpt_manager.load_ckpt(
model.get_network(),
by_score=config.EVAL.load_ckpt_by_score,
name=config.EVAL.ckpt_name)
print(result)
save_path_root = os.path.join(config.EVAL.LOG_DIR.save,
config.EVAL.eval_mode, ckpt_name, date)
exists_or_mkdir(save_path_root)
torch.save(model.get_network().state_dict(),
os.path.join(save_path_root, ckpt_name + '.pytorch'))
##################################################
inp_folder_path_list, _, _ = load_file_list(config.EVAL.inp_path)
ref_folder_path_list, _, _ = load_file_list(config.EVAL.ref_path)
inp_segmap_folder_path_list, _, _ = load_file_list(
config.EVAL.inp_segmap_path)
ref_segmap_folder_path_list, _, _ = load_file_list(
config.EVAL.ref_segmap_path)
print(toGreen('Starting Color Trasfer'))
itr = 0
for folder_idx in np.arange(len(inp_folder_path_list)):
inp_folder_path = inp_folder_path_list[folder_idx]
ref_video_path = ref_folder_path_list[folder_idx]
inp_segmap_folder_path = inp_segmap_folder_path_list[folder_idx]
ref_segmap_video_path = ref_segmap_folder_path_list[folder_idx]
_, inp_file_path_list, _ = load_file_list(inp_folder_path)
_, ref_file_path_list, _ = load_file_list(ref_video_path)
_, inp_segmap_file_path_list, _ = load_file_list(
inp_segmap_folder_path)
_, ref_segmap_file_path_list, _ = load_file_list(ref_segmap_video_path)
for file_idx in np.arange(len(inp_file_path_list)):
inp_path = inp_file_path_list[file_idx]
ref_path = ref_file_path_list[file_idx]
inp_segmap_path = inp_segmap_file_path_list[file_idx]
ref_segmap_path = ref_segmap_file_path_list[file_idx]
# inputs['inp'], inputs['inp_hist'] = torch.FloatTensor(_read_frame_cv(inp_path, config).transpose(0, 3, 1, 2)).cuda()
inputs['inp'], inputs['inp_hist'] = _read_frame_cv(
inp_path, config)
inputs['ref'], inputs['ref_hist'] = _read_frame_cv(
ref_path, config)
inputs['seg_inp'], inputs['seg_ref'] = _read_segmap(
inputs['inp'], inp_segmap_path, ref_segmap_path, config.is_rep)
p = 30
reppad = torch.nn.ReplicationPad2d(p)
for key, val in inputs.items():
inputs[key] = torch.FloatTensor(inputs[key].transpose(
0, 3, 1, 2)).cuda()
inputs['inp'] = reppad(inputs['inp'])
inputs['ref'] = reppad(inputs['ref'])
inputs['seg_inp'] = reppad(inputs['seg_inp'])
inputs['seg_ref'] = reppad(inputs['seg_ref'])
with torch.no_grad():
outs = model.get_results(inputs)
inputs['inp'] = inputs['inp'][:, :, p:(inputs['inp'].size(2) - p),
p:(inputs['inp'].size(3) - p)]
inputs['ref'] = inputs['ref'][:, :, p:(inputs['ref'].size(2) - p),
p:(inputs['ref'].size(3) - p)]
outs['result'] = outs['result'][:, :,
p:(outs['result'].size(2) - p),
p:(outs['result'].size(3) - p)]
outs['result_idt'] = outs['result_idt'][:, :, p:(
outs['result_idt'].size(2) -
p), p:(outs['result_idt'].size(3) - p)]
outs['seg_inp'] = outs['seg_inp'][:, :,
p:(outs['seg_inp'].size(2) - p),
p:(outs['seg_inp'].size(3) - p)]
outs['seg_ref'] = outs['seg_ref'][:, :,
p:(outs['seg_ref'].size(2) - p),
p:(outs['seg_ref'].size(3) - p)]
file_name = os.path.basename(
inp_file_path_list[file_idx]).split('.')[0]
save_path = save_path_root
exists_or_mkdir(save_path)
vutils.save_image(LAB2RGB_cv(inputs['inp'].detach().cpu(),
config.type),
'{}/{}_1_inp.png'.format(save_path, itr),
nrow=3,
padding=0,
normalize=False)
vutils.save_image(LAB2RGB_cv(inputs['ref'].detach().cpu(),
config.type),
'{}/{}_2_ref.png'.format(save_path, itr),
nrow=3,
padding=0,
normalize=False)
i = 3
out_keys = ['result', 'result_idt', 'seg_inp', 'seg_ref']
for key, val in outs.items():
if key in out_keys:
if val is not None:
if 'seg' in key:
if config.identity is False and 'idt' in key:
continue
vutils.save_image(
LAB2RGB_cv(val.detach().cpu(), 'rgb') / 8.,
'{}/{}_{}_out_{}.png'.format(
save_path, itr, i, key),
nrow=3,
padding=0,
normalize=False)
else:
if config.identity is False and 'idt' in key:
continue
vutils.save_image(LAB2RGB_cv(
val.detach().cpu(), config.type),
'{}/{}_{}_out_{}.png'.format(
save_path, itr, i, key),
nrow=3,
padding=0,
normalize=False)
i += 1
#PSNR
print(
'[{}][{}/{}'.format(ckpt_name, folder_idx + 1,
len(inp_folder_path_list)),
'{}/{}]'.format(file_idx + 1, len(inp_file_path_list)))
itr += 1
def main():
if args.seed is None:
args.seed = random.randint(0, 2 ** 31 - 1)
tf.set_random_seed(args.seed)
np.random.seed(args.seed)
random.seed(args.seed)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
if args.mode == "test" or args.mode == "export":
if args.checkpoint is None:
raise Exception("checkpoint required for test mode")
# load some options from the checkpoint
options = {"which_direction", "ngf", "ndf", "lab_colorization"}
with open(os.path.join(args.checkpoint, "options.json")) as f:
for key, val in json.loads(f.read()).items():
if key in options:
print("loaded", key, "=", val)
setattr(args, key, val)
# disable these features in test mode
args.scale_size = CROP_SIZE
args.flip = False
for k, v in args._get_kwargs():
print(k, "=", v)
with open(os.path.join(args.output_dir, "options.json"), "w") as f:
f.write(json.dumps(vars(args), sort_keys=True, indent=4))
if args.mode == "export":
export_model()
return
with tf.device("/cpu:0"):
examples = load_examples()
print("examples count = %d" % examples.count)
# queue = tf.RandomShuffleQueue(20, 5, dtypes=tf.float32)
# enqueue_op = queue.enqueue([examples.inputs, examples.targets])
#
# inputs = queue.dequeue_many(args.batch_size)
# inputs and targets are [batch_size, height, width, channels]
model = create_model(examples.inputs, examples.targets)
# undo colorization splitting on images that we use for display/output
inputs = deprocess(examples.inputs)
targets = deprocess(examples.targets)
outputs = deprocess(model.outputs)
def convert(image):
if args.aspect_ratio != 1.0:
# upscale to correct aspect ratio
size = [CROP_SIZE, int(round(CROP_SIZE * args.aspect_ratio))]
image = tf.image.resize_images(image, size=size, method=tf.image.ResizeMethod.BICUBIC)
return tf.image.convert_image_dtype(image, dtype=tf.uint8, saturate=True)
# reverse any processing on images so they can be written to disk or displayed to user
with tf.name_scope("convert_inputs"):
converted_inputs = convert(inputs)
with tf.name_scope("convert_targets"):
converted_targets = convert(targets)
with tf.name_scope("convert_outputs"):
converted_outputs = convert(outputs)
with tf.name_scope("encode_images"):
display_fetches = {
"paths": examples.paths,
"inputs": tf.map_fn(tf.image.encode_png, converted_inputs, dtype=tf.string, name="input_pngs"),
"targets": tf.map_fn(tf.image.encode_png, converted_targets, dtype=tf.string, name="target_pngs"),
"outputs": tf.map_fn(tf.image.encode_png, converted_outputs, dtype=tf.string, name="output_pngs"),
}
# summaries
with tf.name_scope("inputs_summary"):
tf.summary.image("inputs", converted_inputs)
with tf.name_scope("targets_summary"):
tf.summary.image("targets", converted_targets)
with tf.name_scope("outputs_summary"):
tf.summary.image("outputs", converted_outputs)
with tf.name_scope("predict_real_summary"):
tf.summary.image("predict_real", tf.image.convert_image_dtype(model.predict_real, dtype=tf.uint8))
with tf.name_scope("predict_fake_summary"):
tf.summary.image("predict_fake", tf.image.convert_image_dtype(model.predict_fake, dtype=tf.uint8))
tf.summary.scalar("discriminator_loss", model.discrim_loss)
tf.summary.scalar("generator_loss_GAN", model.gen_loss_GAN)
tf.summary.scalar("generator_loss_L1", model.gen_loss_L1)
if args.summarize_histograms:
for var in tf.trainable_variables():
tf.summary.histogram(var.op.name + "/values", var)
for grad, var in model.discrim_grads_and_vars + model.gen_grads_and_vars:
tf.summary.histogram(var.op.name + "/gradients", grad)
if args.count_parameters:
with tf.name_scope("parameter_count"):
parameter_count = tf.reduce_sum([tf.reduce_prod(tf.shape(v)) for v in tf.trainable_variables()])
saver = tf.train.Saver(max_to_keep=1)
logdir = args.output_dir if (args.trace_freq > 0 or args.summary_freq > 0) else None
sv = tf.train.Supervisor(logdir=logdir, save_summaries_secs=0, saver=None)
with sv.managed_session() as sess:
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
if args.count_parameters:
print("parameter_count =", sess.run(parameter_count))
if args.checkpoint is not None:
print("loading model from checkpoint")
checkpoint = tf.train.latest_checkpoint(args.checkpoint)
saver.restore(sess, checkpoint)
max_steps = 2 ** 32
if args.max_epochs is not None:
max_steps = examples.steps_per_epoch * args.max_epochs
if args.max_steps is not None:
max_steps = args.max_steps
if args.mode == "test":
# testing
# at most, process the test data once
start = time.time()
max_steps = min(examples.steps_per_epoch, max_steps)
for step in range(max_steps):
results = sess.run(display_fetches)
filesets = save_images(results)
for i, f in enumerate(filesets):
print("evaluated image", f["name"])
index_path = append_index(filesets)
print("wrote index at", index_path)
print("rate", (time.time() - start) / max_steps)
else:
# training
start = time.time()
for step in range(max_steps):
def should(freq):
return freq > 0 and ((step + 1) % freq == 0 or step == max_steps - 1)
options = None
run_metadata = None
if should(args.trace_freq):
options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
fetches = {
"train": model.train,
"global_step": sv.global_step,
}
if should(args.progress_freq):
fetches["discrim_loss"] = model.discrim_loss
fetches["gen_loss_GAN"] = model.gen_loss_GAN
fetches["gen_loss_L1"] = model.gen_loss_L1
if should(args.summary_freq):
fetches["summary"] = sv.summary_op
if should(args.display_freq):
fetches["display"] = display_fetches
results = sess.run(fetches, options=options, run_metadata=run_metadata)
if should(args.summary_freq):
print("recording summary")
sv.summary_writer.add_summary(results["summary"], results["global_step"])
if should(args.display_freq):
print("saving display images")
filesets = save_images(results["display"], step=results["global_step"])
append_index(filesets, step=True)
if should(args.trace_freq):
print("recording trace")
sv.summary_writer.add_run_metadata(run_metadata, "step_%d" % results["global_step"])
if should(args.progress_freq):
# global_step will have the correct step count if we resume from a checkpoint
train_epoch = math.ceil(results["global_step"] / examples.steps_per_epoch)
train_step = (results["global_step"] - 1) % examples.steps_per_epoch + 1
rate = (step + 1) * args.batch_size / (time.time() - start)
remaining = (max_steps - step) * args.batch_size / rate
print("progress epoch %d step %d image/sec %0.1f remaining %dm" % (
train_epoch, train_step, rate, remaining / 60))
print("discrim_loss", results["discrim_loss"])
print("gen_loss_GAN", results["gen_loss_GAN"])
print("gen_loss_L1", results["gen_loss_L1"])
if should(args.save_freq):
print("saving model")
saver.save(sess, os.path.join(args.output_dir, "model"), global_step=sv.global_step)
if sv.should_stop():
break
coord.request_stop()
coord.join(threads)
def __init__(self, camera_stream, obstacle_tracking_stream, flags,
camera_setup):
from dataset.dataset_factory import get_dataset
from model.model import create_model, load_model
from opts import opts
from utils.tracker import Tracker
camera_stream.add_callback(self.on_frame_msg,
[obstacle_tracking_stream])
self._flags = flags
self._logger = erdos.utils.setup_logging(self.config.name,
self.config.log_file_name)
self._csv_logger = erdos.utils.setup_csv_logging(
self.config.name + '-csv', self.config.csv_log_file_name)
self._camera_setup = camera_setup
# TODO(ionel): Might have to filter labels when running with a coco
# and a nuscenes model.
num_classes = {
'kitti_tracking': 3,
'coco': 90,
'mot': 1,
'nuscenes': 10
}
# Other flags:
# 1) --K ; max number of output objects.
# 2) --fix_short ; resizes the height of the image to fix short, and
# the width such the aspect ratio is maintained.
# 3) --pre_hm ; pre heat map.
# 4) --input_w; str(camera_setup.width)
# 5) --input_h; str(camera_setup.height)
args = [
'tracking', '--load_model', flags.center_track_model_path,
'--dataset', flags.center_track_model, '--test_focal_length',
str(int(camera_setup.get_focal_length())), '--out_thresh',
str(flags.obstacle_detection_min_score_threshold), '--pre_thresh',
str(flags.obstacle_detection_min_score_threshold), '--new_thresh',
str(flags.obstacle_detection_min_score_threshold),
'--track_thresh',
str(flags.obstacle_detection_min_score_threshold), '--max_age',
str(flags.obstacle_track_max_age), '--num_classes',
str(num_classes[flags.center_track_model]), '--tracking',
'--hungarian'
]
opt = opts().init(args)
gpu = True
if gpu:
opt.device = torch.device('cuda')
else:
opt.device = torch.device('cpu')
self.opt = opt
self.model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
self.model = load_model(self.model, opt.load_model, opt)
self.model = self.model.to(self.opt.device)
self.model.eval()
self.trained_dataset = get_dataset(opt.dataset)
self.mean = np.array(self.trained_dataset.mean,
dtype=np.float32).reshape(1, 1, 3)
self.std = np.array(self.trained_dataset.std,
dtype=np.float32).reshape(1, 1, 3)
self.rest_focal_length = self.trained_dataset.rest_focal_length \
if self.opt.test_focal_length < 0 else self.opt.test_focal_length
self.flip_idx = self.trained_dataset.flip_idx
self.cnt = 0
self.pre_images = None
self.pre_image_ori = None
self.tracker = Tracker(opt)
def main(opt):
torch.manual_seed(opt.seed)
torch.backends.cudnn.benchmark = not opt.not_cuda_benchmark and not opt.eval
Dataset = get_dataset(opt.dataset)
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
print(opt)
if not opt.not_set_cuda_env:
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
opt.device = torch.device('cuda' if opt.gpus[0] >= 0 else 'cpu')
logger = Logger(opt)
print('Creating model...')
model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
optimizer = get_optimizer(opt, model)
start_epoch = 0
lr = opt.lr
if opt.load_model != '':
model, optimizer, start_epoch = load_model(
model, opt.load_model, opt, optimizer)
trainer = Trainer(opt, model, optimizer)
trainer.set_device(opt.gpus, opt.chunk_sizes, opt.device)
if opt.val_intervals < opt.num_epochs or opt.eval:
print('Setting up validation data...')
val_loader = torch.utils.data.DataLoader(
Dataset(opt, opt.val_split), batch_size=1, shuffle=False,
num_workers=1, pin_memory=True)
if opt.eval:
_, preds = trainer.val(0, val_loader)
val_loader.dataset.run_eval(preds, opt.save_dir, n_plots=opt.eval_n_plots,
render_curves=opt.eval_render_curves)
return
print('Setting up train data...')
train_loader = torch.utils.data.DataLoader(
Dataset(opt, opt.train_split), batch_size=opt.batch_size,
shuffle=opt.shuffle_train, num_workers=opt.num_workers,
pin_memory=True, drop_last=True
)
print('Starting training...')
for epoch in range(start_epoch + 1, opt.num_epochs + 1):
mark = epoch if opt.save_all else 'last'
# log learning rate
for param_group in optimizer.param_groups:
lr = param_group['lr']
logger.scalar_summary('LR', lr, epoch)
break
# train one epoch
log_dict_train, _ = trainer.train(epoch, train_loader)
logger.write('epoch: {} |'.format(epoch))
# log train results
for k, v in log_dict_train.items():
logger.scalar_summary('train_{}'.format(k), v, epoch)
logger.write('{} {:8f} | '.format(k, v))
# evaluate
if opt.val_intervals > 0 and epoch % opt.val_intervals == 0:
save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(mark)),
epoch, model, optimizer)
with torch.no_grad():
log_dict_val, preds = trainer.val(epoch, val_loader)
# evaluate val set using dataset-specific evaluator
if opt.run_dataset_eval:
out_dir = val_loader.dataset.run_eval(preds, opt.save_dir,
n_plots=opt.eval_n_plots,
render_curves=opt.eval_render_curves)
# log dataset-specific evaluation metrics
with open('{}/metrics_summary.json'.format(out_dir), 'r') as f:
metrics = json.load(f)
logger.scalar_summary('AP/overall', metrics['mean_ap']*100.0, epoch)
for k,v in metrics['mean_dist_aps'].items():
logger.scalar_summary('AP/{}'.format(k), v*100.0, epoch)
for k,v in metrics['tp_errors'].items():
logger.scalar_summary('Scores/{}'.format(k), v, epoch)
logger.scalar_summary('Scores/NDS', metrics['nd_score'], epoch)
# log eval results
for k, v in log_dict_val.items():
logger.scalar_summary('val_{}'.format(k), v, epoch)
logger.write('{} {:8f} | '.format(k, v))
# save this checkpoint
else:
save_model(os.path.join(opt.save_dir, 'model_last.pth'),
epoch, model, optimizer)
logger.write('\n')
if epoch in opt.save_point:
save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)),
epoch, model, optimizer)
# update learning rate
if epoch in opt.lr_step:
lr = opt.lr * (0.1 ** (opt.lr_step.index(epoch) + 1))
print('Drop LR to', lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
logger.close()
import torch
import torch.utils.data
from opts import opts
from model.model import create_model, load_model, save_model
from model.data_parallel import DataParallel
from logger import Logger
from dataset.dataset_factory import get_dataset
from trainer import Trainer
from main import get_optimizer
if __name__ == '__main__':
opt = opts().parse()
torch.manual_seed(opt.seed)
Dataset = get_dataset(opt.dataset)
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
path_1 = '/mnt/3dvision-cpfs/zhuoyu/CenterTrack/exp/ddd/nu_3d_det_uni/model_last.pth'
path_2 = '/mnt/3dvision-cpfs/zhuoyu/CenterTrack/exp/ddd/nu_3d_det_fix_param/model_last.pth'
model_1 = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
model_2 = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
optimizer = get_optimizer(opt, model_1)
model_1, _, _ = load_model(model_1, path_1, opt, optimizer)
model_2, _, _ = load_model(model_2, path_2, opt, optimizer)
for p1, p2 in zip(model_1.parameters(), model_2.parameters()):
if p1.data.ne(p2.data).sum() > 0:
print(False)
else:
print(True)