def main():
data_loader = DataLoader(args)
if (args.process_data):
data_loader.process_data()
return
torch.cuda.set_device(args.gpu_device)
data_loader.load()
if (args.model == "rnn"):
myModel = model.RNNModel(args, data_loader.vocab_size, 8,
data_loader.id_2_vec).cuda()
elif (args.model == "cnn"):
myModel = model.CNNModel(args, data_loader.vocab_size, 8,
data_loader.id_2_vec).cuda()
elif (args.model == "baseline"):
myModel = model.Baseline(args, data_loader.vocab_size, 8,
data_loader.id_2_vec).cuda()
else:
print("invalid model type")
exit(1)
if (args.test_only):
test(myModel, data_loader, args)
else:
train(myModel, data_loader, args)
def main(objects, **kwargs):
nets = [
model.Net(objects).cuda(),
model.Baseline(objects).cuda(),
]
loader = get_loader(objects, **kwargs)
plins = run(nets, loader, 1000, train=True)
accs = run(nets, loader, 200, train=False)
return {'plins': plins, 'accs': accs}
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-l',
'--load_from',
type=str,
default='models/deblur.hdf5')
parser.add_argument('-s',
'--save_to',
type=str,
default='models/deblur.tflite')
parser.add_argument('-d', '--depth', type=int, default=4)
parser.add_argument('-t', '--test', type=str, default='example/input.png')
parser.add_argument('-o', '--optimize', type=str, default='')
parser.add_argument('-q', '--quantize', type=str, default='')
cfg = parser.parse_args()
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
tf.config.experimental.set_memory_growth(gpus[0], True)
except RuntimeError as e:
print(e)
# Prepare the test input
test_input = imageio.imread(cfg.test)
test_input = data.normalize(test_input)
test_input = np.expand_dims(test_input, axis=0)
_, h, w, c = test_input.shape
representative = 'REDS/{}/train_blur'
if h == 256 and w == 256:
print('hi')
representative = representative.format('train_crop')
else:
representative = representative.format('train')
if cfg.depth == 4:
net = model.Baseline(h, w)
else:
net_class = getattr(model, 'Small{}'.format(cfg.depth))
net = net_class(h, w)
net.build(input_shape=(None, h, w, c))
net.load_weights(cfg.load_from)
# Make a dummy prediction to get the input shape
net.predict(test_input, batch_size=1)
net.summary()
# Convert to the TFLite model
converter = lite.TFLiteConverter.from_keras_model(net)
if cfg.optimize == 'weight':
converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_LATENCY]
elif 'integer' in cfg.quantize:
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# Dataset for tuning
def gen_rep():
list_dir = os.listdir(representative)
list_dir.sort()
for d in tqdm.tqdm(list_dir, ncols=80):
imgs = glob.glob(path.join(representative, d, '*.png'))
img = random.choice(imgs)
x = imageio.imread(img)
hh, ww, _ = x.shape
py = random.randrange(0, hh - h + 1)
px = random.randrange(0, ww - w + 1)
x = x[py:(py + h), px:(px + w)]
x = np.expand_dims(x, axis=0)
x = x.astype(np.float32)
x = x - 128
yield [x]
converter.representative_dataset = gen_rep
if 'full' in cfg.quantize:
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
'''
elif 'fp16' in cfg.quantize:
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
'''
lite_model = converter.convert()
with open(cfg.save_to, 'wb') as f:
f.write(lite_model)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--patch_size', type=int, default=128)
parser.add_argument('--keep_range', action='store_true')
parser.add_argument('--batch_size', type=int, default=16)
parser.add_argument('--depth', type=int, default=4)
parser.add_argument('--epochs', type=int, default=20)
parser.add_argument('--lr', type=float, default=1e-4)
parser.add_argument('--lr_gamma', type=float, default=0.5)
parser.add_argument('--milestones', nargs='+', default=[10, 15])
parser.add_argument('--exp_name', type=str, default='baseline')
parser.add_argument('--save_as', type=str, default='models/deblur.hdf5')
cfg = parser.parse_args()
# For checking the GPU usage
#tf.debugging.set_log_device_placement(True)
# For limiting the GPU usage
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
tf.config.experimental.set_memory_growth(gpus[0], True)
except RuntimeError as e:
print(e)
dataset_train = data.REDS(
cfg.batch_size,
patch_size=cfg.patch_size,
train=True,
keep_range=cfg.keep_range,
)
# We have 3,000 validation frames.
# Note that each frame will be center-cropped for the validation.
dataset_val = data.REDS(
20,
patch_size=cfg.patch_size,
train=False,
keep_range=cfg.keep_range,
)
if cfg.depth == 4:
net = model.Baseline(cfg.patch_size, cfg.patch_size)
else:
net_class = getattr(model, 'Small{}'.format(cfg.depth))
net = net_class(cfg.patch_size, cfg.patch_size)
net.build(input_shape=(None, cfg.patch_size, cfg.patch_size, 3))
kwargs = {'optimizer': 'adam', 'loss': 'mse'}
if cfg.keep_range:
net.compile(**kwargs, metrics=[metric.psnr_full])
else:
net.compile(**kwargs, metrics=[metric.psnr])
net.summary()
# Callback functions
# For TensorBoard logging
logging = callbacks.TensorBoard(
log_dir=path.join('logs', cfg.exp_name),
update_freq=100,
)
# For checkpointing
os.makedirs(path.dirname(cfg.save_as), exist_ok=True)
checkpointing = callbacks.ModelCheckpoint(
cfg.save_as,
verbose=1,
save_weights_only=True,
)
def scheduler(epoch):
idx = bisect.bisect_right(cfg.milestones, epoch)
lr = cfg.lr * (cfg.lr_gamma**idx)
return lr
# For learning rate scheduling
scheduling = callbacks.LearningRateScheduler(scheduler, verbose=1)
net.fit_generator(
dataset_train,
epochs=cfg.epochs,
callbacks=[logging, checkpointing, scheduling],
validation_data=dataset_val,
validation_freq=1,
max_queue_size=16,
workers=8,
use_multiprocessing=True,
)