def load_model(model_name, weight_path, input_size, framework):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
if framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
model = None
raise ValueError
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
print('load yolo tiny 3')
elif model_name == 'yolov3':
utils.load_weights_v3(model, weight_path)
print('load yolo 3')
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
print('load yolo 4')
else:
raise ValueError
else:
model.load_weights(weight_path).expect_partial()
print('Restoring weights from: %s ' % weight_path)
return model
def transfer_tflite(model_name, weight_path, output, input_size):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
model = None
raise ValueError
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name == ' yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path).expect_partial()
print('Restoring weights from: %s ... ' % weight_path)
converter = tf.lite.TFLiteConverter.from_keras_model(model)
# converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
tflite_model = converter.convert()
open(output, 'wb').write(tflite_model)
def detect(model_name, weight_path, input_size, image_path, framework):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
if model_name == 'yolov3_tiny':
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, True)
elif model_name == 'yolov3':
STRIDES = np.array(cfg.YOLO.STRIDES)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, False)
elif model_name == 'yolov4':
STRIDES = np.array(cfg.YOLO.STRIDES)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, False)
else:
raise ValueError
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preprocess(np.copy(original_image),
[input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
if framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
model = None
raise ValueError
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
# utils.extract_weights_tiny(model, weight_path)
print('load yolo tiny 3')
elif model_name == 'yolov3':
utils.load_weights_v3(model, weight_path)
print('load yolo 3')
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
print('load yolo 4')
else:
raise ValueError
elif weight_path.split(".")[-1] == "npy":
if model_name == 'yolov3_tiny':
# utils.load_weights_tiny_npy(model, weight_path)
print('load yolo tiny 3 npy')
else:
model.load_weights(weight_path)
print('Restoring weights from: %s ' % weight_path)
# weight = np.load('D:\\coursera\\YoLoSerirs\\checkpoint\\yolo3_tiny.npy', allow_pickle=True)
# model.set_weights(weight)
# model.summary()
start_time = time.time()
pred_bbox = model.predict(image_data)
print(time.time() - start_time)
else:
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path=weight_path)
interpreter.allocate_tensors()
# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
print(output_details)
interpreter.set_tensor(input_details[0]['index'], image_data)
start_time = time.time()
interpreter.invoke()
pred_bbox = [
interpreter.get_tensor(output_details[i]['index'])
for i in range(len(output_details))
]
print(time.time() - start_time)
if model_name == 'yolov4':
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES,
XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size,
input_size, 0.5)
bboxes = utils.nms(bboxes, 0.3, method='nms')
image = visualize.draw_bbox(original_image, bboxes)
image = Image.fromarray(image)
image.show()
def train(model_name, weight_path, save_path, logdir=None):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
trainset = Dataset('train')
testset = Dataset('test')
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
if model_name=='yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name=='yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name=='yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
raise ValueError
# for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']:
# layer = model.get_layer(name)
# print(layer.name, layer.output_shape)
if weight_path:
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name=='yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name=='yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path)
print('Restoring weights from: %s ... ' % weight_path)
optimizer = tf.keras.optimizers.Adam()
if logdir:
if os.path.exists(logdir):
shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)
else:
writer = None
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(pred, conv, target[i][0], target[i][1],
STRIDES=STRIDES, NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
tf.print("=> STEP %4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, optimizer.lr.numpy(),
giou_loss, conf_loss,
prob_loss, total_loss))
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + \
0.5*(cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * \
((1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi)))
optimizer.lr.assign(lr.numpy())
# if writer:
# # writing summary data
# with writer.as_default():
# tf.summary.scalar("lr", optimizer.lr, step=global_steps)
# tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
# tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps)
# tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
# tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
# writer.flush()
def test_step(image_data, target):
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(pred, conv, target[i][0], target[i][1],
STRIDES=STRIDES, NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss, conf_loss,
prob_loss, total_loss))
for epoch in range(first_stage_epochs + second_stage_epochs):
if epoch < first_stage_epochs:
if not isfreeze:
isfreeze = True
for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']:
freeze = model.get_layer(name)
ops.freeze_all(freeze)
elif epoch >= first_stage_epochs:
if isfreeze:
isfreeze = False
for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']:
freeze = model.get_layer(name)
ops.unfreeze_all(freeze)
for image_data, target in trainset:
train_step(image_data, target)
for image_data, target in testset:
test_step(image_data, target)
if save_path:
model.save_weights(save_path)
def run(config):
# 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)
config['resolution'] = utils.imsize_dict[config['dataset']]
config['n_classes'] = utils.nclass_dict[config['dataset']]
config['G_activation'] = utils.activation_dict[config['G_nl']]
config['D_activation'] = utils.activation_dict[config['D_nl']]
# By default, skip init if resuming training.
if config['resume']:
print('Skipping initialization for training resumption...')
config['skip_init'] = True
config = vae_utils.update_config_roots(config)
device = 'cuda'
# Seed RNG
utils.seed_rng(config['seed'])
# Prepare root folders if necessary
utils.prepare_root(config)
# Setup cudnn.benchmark for free speed
torch.backends.cudnn.benchmark = True
# Import the model--this line allows us to dynamically select different files.
model = import_module('Network.' + config['model'])
experiment_name = (config['experiment_name'] if config['experiment_name']
else utils.name_from_config(config))
print('Experiment name is %s' % experiment_name)
# Next, build the model
G = model.Generator(**config).to(device)
D = model.Discriminator(**config).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
# 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()
# Consider automatically reducing SN_eps?
GD = model.G_D(G, D)
print(G)
print(D)
print('Number of params in G: {} D: {}'.format(
*
[sum([p.data.nelement() for p in net.parameters()])
for net in [G, D]]))
# 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,
'config': config
}
# If loading from a pre-trained model, load weights
if config['resume']:
print('Loading weights...')
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)
# If parallel, parallelize the GD module
if config['parallel']:
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.
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=(not config['resume']))
print('Training Metrics will be saved to {}'.format(train_metrics_fname))
train_log = utils.MyLogger(train_metrics_fname,
reinitialize=(not config['resume']),
logstyle=config['logstyle'])
# Write metadata
utils.write_metadata(config['logs_root'], experiment_name, config,
state_dict)
# Prepare data; the Discriminator's batch size is all that needs to be passed
# to the dataloader, as G doesn't require dataloading.
# Note that at every loader iteration we pass in enough data to complete
# a full D iteration (regardless of number of D steps and accumulations)
D_batch_size = (config['batch_size'] * config['num_D_steps'] *
config['num_D_accumulations'])
loaders = vae_utils.get_minidata_loaders(**{
**config, 'batch_size': D_batch_size,
'start_itr': state_dict['itr']
})
# Prepare inception metrics: FID and IS
get_inception_metrics = inception_utils.prepare_inception_metrics(
config['dataset'], config['parallel'], config['data_root'],
config['no_fid'])
# Prepare noise and randomly sampled label arrays
# Allow for different batch sizes in G
G_batch_size = max(config['G_batch_size'], config['batch_size'])
z_, y_ = utils.prepare_z_y(G_batch_size,
G.dim_z,
config['n_classes'],
device=device,
fp16=config['G_fp16'])
# Prepare a fixed z & y to see individual sample evolution throghout training
fixed_z, fixed_y = utils.prepare_z_y(G_batch_size,
G.dim_z,
config['n_classes'],
device=device,
fp16=config['G_fp16'])
fixed_z.sample_()
fixed_y.sample_()
# Loaders are loaded, prepare the training function
if config['which_train_fn'] == 'GAN':
train = train_fns.GAN_training_function(G, D, GD, z_, y_, ema,
state_dict, config)
# Else, assume debugging and use the dummy train fn
else:
train = train_fns.dummy_training_function()
# 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),
z_=z_,
y_=y_,
config=config)
print('Beginning training at epoch %d...' % state_dict['epoch'])
# Train for specified number of epochs, although we mostly track G iterations.
for epoch in range(state_dict['epoch'], config['num_epochs']):
# Which progressbar to use? TQDM or my own?
if config['pbar'] == 'mine':
pbar = utils.progress(loaders[0],
displaytype='s1k' if
config['use_multiepoch_sampler'] else 'eta')
else:
pbar = tqdm(loaders[0])
for i, (x, y) in enumerate(pbar):
# 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()
if config['D_fp16']:
x, y = x.to(device).half(), y.to(device)
else:
x, y = x.to(device), y.to(device)
metrics = train(x, y)
train_log.log(itr=int(state_dict['itr']), **metrics)
# Every sv_log_interval, log singular values
if (config['sv_log_interval'] > 0) and (
not (state_dict['itr'] % config['sv_log_interval'])):
train_log.log(itr=int(state_dict['itr']),
**{
**utils.get_SVs(G, 'G'),
**utils.get_SVs(D, 'D')
})
# If using my progbar, print metrics.
if config['pbar'] == 'mine':
print(', '.join(
['itr: %d' % state_dict['itr']] +
['%s : %+4.3f' % (key, metrics[key]) for key in metrics]),
end=' ')
# Save weights and copies as configured at specified interval
if not (state_dict['itr'] % config['save_every']):
if config['G_eval_mode']:
print('Switchin G to eval mode...')
G.eval()
if config['ema']:
G_ema.eval()
train_fns.save_and_sample(G, D, G_ema, z_, y_, fixed_z,
fixed_y, state_dict, config,
experiment_name)
# Test every specified interval
if not (state_dict['itr'] % config['test_every']):
if config['G_eval_mode']:
print('Switchin G to eval mode...')
G.eval()
train_fns.test(G, D, G_ema, z_, y_, state_dict, config, sample,
get_inception_metrics, experiment_name,
test_log)
# Increment epoch counter at end of epoch
state_dict['epoch'] += 1
def run(config):
# 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,
'config': config
}
# Optionally, get the configuration from the state dict. This allows for
# recovery of the config provided only a state dict and experiment name,
# and can be convenient for writing less verbose sample shell scripts.
if config['config_from_name']:
utils.load_weights(None,
None,
state_dict,
config['weights_root'],
config['experiment_name'],
config['load_weights'],
None,
strict=False,
load_optim=False)
# Ignore items which we might want to overwrite from the command line
for item in state_dict['config']:
if item not in [
'z_var', 'base_root', 'batch_size', 'G_batch_size',
'use_ema', 'G_eval_mode'
]:
config[item] = state_dict['config'][item]
# update config (see train.py for explanation)
config['resolution'] = utils.imsize_dict[config['dataset']]
config['n_classes'] = utils.nclass_dict[config['dataset']]
config['G_activation'] = utils.activation_dict[config['G_nl']]
config['D_activation'] = utils.activation_dict[config['D_nl']]
config = utils.update_config_roots(config)
config['skip_init'] = True
config['no_optim'] = True
device = 'cuda'
# Seed RNG
utils.seed_rng(config['seed'])
# Setup cudnn.benchmark for free speed
torch.backends.cudnn.benchmark = 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)
G = model.Generator(**config).cuda()
utils.count_parameters(G)
# Load weights
print('Loading weights...')
# Here is where we deal with the ema--load ema weights or load normal weights
utils.load_weights(G if not (config['use_ema']) else None,
None,
state_dict,
config['weights_root'],
experiment_name,
config['load_weights'],
G if config['ema'] and config['use_ema'] else None,
strict=False,
load_optim=False)
# Update batch size setting used for G
G_batch_size = max(config['G_batch_size'], config['batch_size'])
z_, y_ = utils.prepare_z_y(G_batch_size,
G.dim_z,
config['n_classes'],
device=device,
fp16=config['G_fp16'],
z_var=config['z_var'])
if config['G_eval_mode']:
print('Putting G in eval mode..')
G.eval()
else:
print('G is in %s mode...' % ('training' if G.training else 'eval'))
#Sample function
sample = functools.partial(utils.sample, G=G, z_=z_, y_=y_, config=config)
if config['accumulate_stats']:
print('Accumulating standing stats across %d accumulations...' %
config['num_standing_accumulations'])
utils.accumulate_standing_stats(G, z_, y_, config['n_classes'],
config['num_standing_accumulations'])
# Sample a number of images and save them to an NPZ, for use with TF-Inception
if config['sample_npz']:
# Lists to hold images and labels for images
x, y = [], []
print('Sampling %d images and saving them to npz...' %
config['sample_num_npz'])
for i in trange(
int(np.ceil(config['sample_num_npz'] / float(G_batch_size)))):
with torch.no_grad():
images, labels = sample()
x += [np.uint8(255 * (images.cpu().numpy() + 1) / 2.)]
y += [labels.cpu().numpy()]
x = np.concatenate(x, 0)[:config['sample_num_npz']]
y = np.concatenate(y, 0)[:config['sample_num_npz']]
print('Images shape: %s, Labels shape: %s' % (x.shape, y.shape))
npz_filename = '%s/%s/samples.npz' % (config['samples_root'],
experiment_name)
print('Saving npz to %s...' % npz_filename)
np.savez(npz_filename, **{'x': x, 'y': y})
# Prepare sample sheets
if config['sample_sheets']:
print('Preparing conditional sample sheets...')
utils.sample_sheet(
G,
classes_per_sheet=utils.classes_per_sheet_dict[config['dataset']],
num_classes=config['n_classes'],
samples_per_class=10,
parallel=config['parallel'],
samples_root=config['samples_root'],
experiment_name=experiment_name,
folder_number=config['sample_sheet_folder_num'],
z_=z_,
)
# Sample interp sheets
if config['sample_interps']:
print('Preparing interp sheets...')
for fix_z, fix_y in zip([False, False, True], [False, True, False]):
utils.interp_sheet(G,
num_per_sheet=16,
num_midpoints=8,
num_classes=config['n_classes'],
parallel=config['parallel'],
samples_root=config['samples_root'],
experiment_name=experiment_name,
folder_number=config['sample_sheet_folder_num'],
sheet_number=0,
fix_z=fix_z,
fix_y=fix_y,
device='cuda')
# Sample random sheet
if config['sample_random']:
print('Preparing random sample sheet...')
images, labels = sample()
torchvision.utils.save_image(images.float(),
'%s/%s/random_samples.jpg' %
(config['samples_root'], experiment_name),
nrow=int(G_batch_size**0.5),
normalize=True)
# Get Inception Score and FID
get_inception_metrics = inception_utils.prepare_inception_metrics(
config['dataset'], config['parallel'], config['no_fid'])
# Prepare a simple function get metrics that we use for trunc curves
def get_metrics():
sample = functools.partial(utils.sample,
G=G,
z_=z_,
y_=y_,
config=config)
IS_mean, IS_std, FID = get_inception_metrics(
sample,
config['num_inception_images'],
num_splits=10,
prints=False)
# Prepare output string
outstring = 'Using %s weights ' % ('ema'
if config['use_ema'] else 'non-ema')
outstring += 'in %s mode, ' % ('eval' if config['G_eval_mode'] else
'training')
outstring += 'with noise variance %3.3f, ' % z_.var
outstring += 'over %d images, ' % config['num_inception_images']
if config['accumulate_stats'] or not config['G_eval_mode']:
outstring += 'with batch size %d, ' % G_batch_size
if config['accumulate_stats']:
outstring += 'using %d standing stat accumulations, ' % config[
'num_standing_accumulations']
outstring += 'Itr %d: PYTORCH UNOFFICIAL Inception Score is %3.3f +/- %3.3f, PYTORCH UNOFFICIAL FID is %5.4f' % (
state_dict['itr'], IS_mean, IS_std, FID)
print(outstring)
if config['sample_inception_metrics']:
print('Calculating Inception metrics...')
get_metrics()
# Sample truncation curve stuff. This is basically the same as the inception metrics code
if config['sample_trunc_curves']:
start, step, end = [
float(item) for item in config['sample_trunc_curves'].split('_')
]
print(
'Getting truncation values for variance in range (%3.3f:%3.3f:%3.3f)...'
% (start, step, end))
for var in np.arange(start, end + step, step):
z_.var = var
# Optionally comment this out if you want to run with standing stats
# accumulated at one z variance setting
if config['accumulate_stats']:
utils.accumulate_standing_stats(
G, z_, y_, config['n_classes'],
config['num_standing_accumulations'])
get_metrics()
def evaluate(model_name, weight_path):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
trainset = Dataset('train')
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
if model_name=='yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name=='yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name=='yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
raise ValueError
if weight_path:
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name=='yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name=='yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path)
print('Restoring weights from: %s ... ' % weight_path)
trainset = Dataset('train')
for image_data, target in trainset:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(pred, conv, target[i][0], target[i][1],
STRIDES=STRIDES, NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
tf.print("=> STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss,
conf_loss, prob_loss, total_loss))
def prune_train(model_name, weight_path, logdir, save_path, epoches):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
trainset = Dataset('train')
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
input_layer = tf.keras.layers.Input(
[cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i,
XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
raise ValueError
if weight_path:
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name == 'yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path)
print('Restoring weights from: %s ... ' % weight_path)
optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)
if os.path.exists(logdir):
shutil.rmtree(logdir)
# for layer in model.layers:
# print(layer.name, isinstance(layer, tf.keras.layers.Conv2D))
def apply_pruning_to_dense(layer):
if isinstance(layer, tf.keras.layers.Conv2D):
return tfmot.sparsity.keras.prune_low_magnitude(layer)
return layer
# Use `tf.keras.models.clone_model` to apply `apply_pruning_to_dense`
# to the layers of the model.
model_for_pruning = tf.keras.models.clone_model(
model,
clone_function=apply_pruning_to_dense,
)
# model_for_pruning.summary()
unused_arg = -1
model_for_pruning.optimizer = optimizer
step_callback = tfmot.sparsity.keras.UpdatePruningStep()
step_callback.set_model(model_for_pruning)
log_callback = tfmot.sparsity.keras.PruningSummaries(
log_dir=logdir) # Log sparsity and other metrics in Tensorboard.
log_callback.set_model(model_for_pruning)
step_callback.on_train_begin() # run pruning callback
for epoch in range(epoches):
log_callback.on_epoch_begin(epoch=unused_arg) # run pruning callback
for image_data, target in trainset:
step_callback.on_train_batch_begin(
batch=unused_arg) # run pruning callback
with tf.GradientTape() as tape:
pred_result = model_for_pruning(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(
pred,
conv,
target[i][0],
target[i][1],
STRIDES=STRIDES,
NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH,
i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(
total_loss, model_for_pruning.trainable_variables)
optimizer.apply_gradients(
zip(gradients, model_for_pruning.trainable_variables))
tf.print(
"=> STEP %4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" %
(global_steps, optimizer.lr.numpy(), giou_loss, conf_loss,
prob_loss, total_loss))
step_callback.on_epoch_end(batch=unused_arg) # run pruning callback
model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)
return model_for_export
def save_tflite(model_name, weight_path, quantize_mode, output, input_size):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
assert quantize_mode in ['int8', 'float16', 'full_int8']
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
model = None
raise ValueError
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name == ' yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path).expect_partial()
print('Restoring weights from: %s ... ' % weight_path)
# model.summary()
converter = tf.lite.TFLiteConverter.from_keras_model(model)
if tf.__version__ >= '2.2.0':
converter.experimental_new_converter = False
if quantize_mode == 'int8':
converter.optimizations = [tf.lite.Optimize.DEFAULT]
elif quantize_mode == 'float16':
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [
tf.compat.v1.lite.constants.FLOAT16
]
elif quantize_mode == 'full_int8':
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
converter.allow_custom_ops = True
converter.representative_dataset = representative_data_gen
else:
raise ValueError
tflite_model = converter.convert()
open(output, 'wb').write(tflite_model)
logging.info("model saved to: {}".format(output))