def load_model(self):
"""
加载权重
"""
if self.cuda:
print("=====> use gpu id: '{}'".format(self.gpu_number))
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu_number
if not torch.cuda.is_available():
raise Exception(
"no GPU found or wrong gpu id, please run without --cuda")
# build the model
self.model = build_model(self.net_type, num_classes=self.classes)
if self.cuda:
self.model = self.model.cuda() # using GPU for inference
cudnn.benchmark = True
if os.path.isfile(self.model_path):
print(f"=====> loading checkpoint '{self.model_path}'")
checkpoint = torch.load(self.model_path)
self.model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
else:
print("=====> no checkpoint found at '{self.model_path}'")
raise FileNotFoundError(
f"no checkpoint found at '{self.model_path}'")
if not os.path.exists(self.save_seg_dir):
os.makedirs(self.save_seg_dir)
def freeze_graph(args, output_node_names):
if not tf.gfile.Exists(args.model_dir):
raise AssertionError(
"Export directory doesn't exists. Please specify an export "
"directory: %s" % args.model_dir)
if not output_node_names:
print("You need to supply the name of a node to --output_node_names.")
return -1
# We retrieve our checkpoint fullpath
# checkpoint = tf.train.get_checkpoint_state(model_dir)
# input_checkpoint = checkpoint.model_checkpoint_path
# We precise the file fullname of our freezed graph
absolute_model_dir = "/".join(args.model_dir.split('/')[:-1])
output_graph = os.path.join(absolute_model_dir, "unfrozen_model_softmax_output.pb")
with tf.Graph().as_default() as graph:
net_input = tf.placeholder(tf.float32, shape=[1, args.crop_height, args.crop_width, 3]) # NHWC format
net, init_fn = model_builder.build_model(model_name=args.model, frontend=args.frontend, net_input=net_input,
num_classes=num_classes, crop_width=args.crop_width,
crop_height=args.crop_height, is_training=False)
net = tf.nn.softmax(net, name="output_name_softmax")
net = tf.argmax(net, name="output_name_argmax", axis=3)
graph_def = graph.as_graph_def()
with gfile.GFile(output_graph, 'wb') as f:
f.write(graph_def.SerializeToString())
print('Successfull written to', output_graph)
def __init__(self, modelName, width, height, checkpoints):
self.model = modelName
self.width = width
self.height = height
self.checkpoints = checkpoints
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
self.class_names_list, self.label_values = helpers.get_label_info(
"class_dict.csv")
self.num_classes = len(self.label_values)
self.config = tf.ConfigProto()
self.config.gpu_options.allow_growth = True
self.config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
self.sess = tf.Session(config=self.config)
self.net_input = tf.placeholder(tf.float32,
shape=[None, None, None, 3])
self.net_output = tf.placeholder(
tf.float32, shape=[None, None, None, self.num_classes])
self.network, _ = model_builder.build_model(
modelName,
net_input=self.net_input,
num_classes=self.num_classes,
crop_width=width,
crop_height=height,
is_training=False)
self.sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
self.saver = tf.train.Saver(max_to_keep=1000)
self.saver.restore(self.sess, checkpoints)
print("Setup done!")
def main(args):
"""
main function for testing
param args: global arguments
return: None
"""
t = PrettyTable(['args_name', 'args_value'])
for k in list(vars(args).keys()):
t.add_row([k, vars(args)[k]])
print(t.get_string(title="Predict Arguments"))
# build the model
model = build_model(args.model, args.classes, args.backbone, args.pretrained, args.out_stride, args.mult_grid)
# load the test set
if args.predict_type == 'validation':
testdataset, class_dict_df = build_dataset_test(args.root, args.dataset, args.crop_size,
mode=args.predict_mode, gt=True)
else:
testdataset, class_dict_df = build_dataset_test(args.root, args.dataset, args.crop_size,
mode=args.predict_mode, gt=False)
DataLoader = data.DataLoader(testdataset, batch_size=args.batch_size,
shuffle=False, num_workers=args.batch_size, pin_memory=True, drop_last=False)
if args.cuda:
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
model = model.cuda()
cudnn.benchmark = True
if not torch.cuda.is_available():
raise Exception("no GPU found or wrong gpu id, please run without --cuda")
if not os.path.exists(args.save_seg_dir):
os.makedirs(args.save_seg_dir)
if args.checkpoint:
if os.path.isfile(args.checkpoint):
checkpoint = torch.load(args.checkpoint)['model']
check_list = [i for i in checkpoint.items()]
# Read weights with multiple cards, and continue training with a single card this time
if 'module.' in check_list[0][0]: # 读取使用多卡训练权重,并且此次使用单卡预测
new_stat_dict = {}
for k, v in checkpoint.items():
new_stat_dict[k[7:]] = v
model.load_state_dict(new_stat_dict, strict=True)
# Read the training weight of a single card, and continue training with a single card this time
else:
model.load_state_dict(checkpoint)
else:
print("no checkpoint found at '{}'".format(args.checkpoint))
raise FileNotFoundError("no checkpoint found at '{}'".format(args.checkpoint))
# define loss function
criterion = build_loss(args, None, 255)
print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n"
">>>>>>>>>>> beginning testing >>>>>>>>>>>>\n"
">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
predict_multiscale_sliding(args=args, model=model, testLoader=DataLoader, class_dict_df=class_dict_df,
scales=args.scales, overlap=args.overlap, criterion=criterion,
mode=args.predict_type, save_result=True)
def main(_):
if not FLAGS.output_file:
raise ValueError('You must supply the path to save to with --output_file')
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default() as graph:
shape = INPUT_SIZE.split(',')
shape = (int(shape[0]), int(shape[1]), 3)
x = tf.placeholder(name='input', dtype = tf.float32, shape = (None, shape[0], shape[1], 3))
img_tf = tf.cast(x, dtype = tf.float32)
# Extract mean.
img_tf -= IMG_MEAN
print(img_tf)
# Create network.
net_input = tf.placeholder(name='input', dtype=tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32, shape=[None, None, None, 32])
net, _ = model_builder.build_model(model_name="MobileUNet", frontend="MobileNetV2", net_input=net_input,
num_classes=17, crop_width=512,
crop_height=512, is_training=False)
# net =
#net = psp_net({'inputs': img_tf}, is_training = False, num_classes = NUM_CLASSES)
# net = unext(img_tf, is_train = False, n_out = NUM_CLASSES)
#net = ICNet_BN({'data': img_tf}, is_training = False, num_classes = NUM_CLASSES)
raw_output = graph.get_tensor_by_name('logits/BiasAdd:0')
# raw_output = net.outputs
#raw_output = net.layers['conv6']
# output = tf.image.resize_bilinear(raw_output, tf.shape(img_tf)[1:3,], name = 'raw_output')
#output = tf.nn.softmax(raw_output)
# output = tf.argmax(raw_output, dimension = 3)
# pred = tf.expand_dims(output, dim = 3, name = 'indices')
# Adding additional params to graph. It is necessary also to point them as outputs in graph freeze conversation, otherwise they will be cuted
# tf.constant(label_colours, name = 'label_colours')
# tf.constant(label_names, name = 'label_names')
shape = INPUT_SIZE.split(',')
shape = (int(shape[0]), int(shape[1]), 3)
tf.constant(shape, name = 'input_size')
tf.constant(["indices"], name = "output_name")
graph_def = graph.as_graph_def()
with gfile.GFile(FLAGS.output_file, 'wb') as f:
f.write(graph_def.SerializeToString())
print('Successfull written to', FLAGS.output_file)
def load_model(): # Load model # tf.reset_default_graph() # restore_graph = tf.Graph() num_classes = 2 session_config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False) session_config.gpu_options.allow_growth = True sess = tf.Session(config=session_config) net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3]) # net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes]) network, init_fn = model_builder.build_model(model_name=args.model, net_input=net_input, num_classes=num_classes, crop_width=args.crop_width, crop_height=args.crop_height, is_training=False) return sess, network, net_input
def predict_model(args):
"""
main function for testing
param args: global arguments
return: None
"""
print(args)
if args.cuda:
print("=====> use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"no GPU found or wrong gpu id, please run without --cuda")
# build the model
model = build_model(args.model, num_classes=args.classes)
if args.cuda:
model = model.cuda() # using GPU for inference
cudnn.benchmark = True
if args.checkpoint:
if os.path.isfile(args.checkpoint):
print("=====> loading checkpoint '{}'".format(args.checkpoint))
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
else:
print("=====> no checkpoint found at '{}'".format(args.checkpoint))
raise FileNotFoundError("no checkpoint found at '{}'".format(
args.checkpoint))
if not os.path.exists(args.save_seg_dir):
os.makedirs(args.save_seg_dir)
# load the test set
if args.use_txt_list:
_, testLoader = build_dataset_test(args.dataset,
args.num_workers,
none_gt=True)
else:
_, testLoader = build_dataset_predict(args.image_input_path,
args.dataset,
args.num_workers,
none_gt=True)
print("=====> beginning testing")
print("test set length: ", len(testLoader))
predict(args, testLoader, model)
def predict(checkpoint_path, datadir, resultdir, model):
class_names_list = ['ruler', 'cell']
label_values = [[255, 0, 0], [255, 255, 255]]
num_classes = len(label_values)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32,
shape=[None, None, None, num_classes])
network, _ = model_builder.build_model('encoder-decoder-skip',
net_input=net_input,
num_classes=num_classes,
crop_width=1024,
crop_height=1024,
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver = tf.train.Saver(max_to_keep=1000)
saver.restore(sess, checkpoint_path)
for image in os.listdir(datadir):
print("Testing image " + image)
file_name = os.path.join(datadir, image)
loaded_image = cv2.imread(file_name)
st = time.time()
output_image = sess.run(network, feed_dict={net_input: input_image})
run_time = time.time() - st
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(
output_image, label_values)
file_name = os.path.join(resultdir, image)
cv2.imwrite(file_name,
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
print("Took %i" % run_time)
# print("Finished!")
# print("Wrote image " + "%s_pred.png"%(os.path.join(args.datadir,image)))
def test_model(args):
"""
main function for testing
param args: global arguments
return: None
"""
print(args)
if args.cuda:
print("use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
# print(args.gpus)
# torch.cuda.set_device(0)
if not torch.cuda.is_available():
raise Exception(
"no GPU found or wrong gpu id, please run without --cuda")
# build the model
model = build_model(args.model, num_classes=args.classes)
if args.cuda:
model = model.cuda() # using GPU for inference
cudnn.benchmark = True
if not os.path.exists(args.save_seg_dir):
os.makedirs(args.save_seg_dir)
# load the test set
datas, testLoader = build_dataset_sliding_test(args.dataset,
args.num_workers,
none_gt=True)
if args.checkpoint:
if os.path.isfile(args.checkpoint):
print("loading checkpoint '{}'".format(args.checkpoint))
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
else:
print("no checkpoint found at '{}'".format(args.checkpoint))
raise FileNotFoundError("no checkpoint found at '{}'".format(
args.checkpoint))
print(">>>>>>>>>>>beginning testing>>>>>>>>>>>")
predict_sliding(args,
model.eval(),
image=testLoader,
tile_size=(args.tile_size, args.tile_size),
classes=args.classes)
def initializeNetwork():
class_names_list, label_values = helpers.get_label_info(
os.path.join(args.dataset, "class_dict.csv"))
num_classes = len(label_values)
print("\n***** Begin prediction *****")
print("Dataset -->", args.dataset)
print("Model -->", args.model_segmentation)
print("Crop Height -->", args.crop_height)
print("Crop Width -->", args.crop_width)
print("Num Classes -->", num_classes)
print("Video -->", args.video)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32,
shape=[None, None, None, num_classes])
network, _ = model_builder.build_model(args.model_segmentation,
net_input=net_input,
num_classes=num_classes,
crop_width=args.crop_width,
crop_height=args.crop_height,
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver = tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
return sess, network, net_input, label_values
def main(args, _log):
def data_augmentation(input_image, output_image):
# Data augmentation
input_image, output_image = utils.random_crop(input_image,
output_image,
args["crop_height"],
args["crop_width"])
if args["h_flip"] and random.randint(0, 1):
input_image = cv2.flip(input_image, 1)
output_image = cv2.flip(output_image, 1)
if args["v_flip"] and random.randint(0, 1):
input_image = cv2.flip(input_image, 0)
output_image = cv2.flip(output_image, 0)
if args["brightness"]:
factor = 1.0 + random.uniform(-1.0 * args["brightness"],
args["brightness"])
table = np.array([((i / 255.0) * factor) * 255
for i in np.arange(0, 256)]).astype(np.uint8)
input_image = cv2.LUT(input_image, table)
if args["rotation"]:
angle = random.uniform(-1 * args["rotation"], args["rotation"])
if args["rotation"]:
M = cv2.getRotationMatrix2D(
(input_image.shape[1] // 2, input_image.shape[0] // 2), angle,
1.0)
input_image = cv2.warpAffine(
input_image,
M, (input_image.shape[1], input_image.shape[0]),
flags=cv2.INTER_NEAREST)
output_image = cv2.warpAffine(
output_image,
M, (output_image.shape[1], output_image.shape[0]),
flags=cv2.INTER_NEAREST)
return input_image, output_image
print("Args:", args)
# Get the names of the classes so we can record the evaluation results
class_names_list, label_values = helpers.get_label_info(
os.path.join(args["dataset"], "class_dict.csv"))
class_names_string = ""
for class_name in class_names_list:
if not class_name == class_names_list[-1]:
class_names_string = class_names_string + class_name + ", "
else:
class_names_string = class_names_string + class_name
num_classes = len(label_values)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# Compute your softmax cross entropy loss
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32,
shape=[None, None, None, num_classes])
network, init_fn = model_builder.build_model(
model_name=args["model"],
frontend=args["frontend"],
net_input=net_input,
num_classes=num_classes,
crop_width=args["crop_width"],
crop_height=args["crop_height"],
is_training=True)
loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(logits=network,
labels=net_output))
opt = tf.train.RMSPropOptimizer(
learning_rate=0.0001, decay=0.995).minimize(
loss, var_list=[var for var in tf.trainable_variables()])
saver = tf.train.Saver(max_to_keep=1000)
sess.run(tf.global_variables_initializer())
utils.count_params()
# If a pre-trained ResNet is required, load the weights.
# This must be done AFTER the variables are initialized with sess.run(tf.global_variables_initializer())
if init_fn is not None:
init_fn(sess)
# Load a previous checkpoint if desired
model_checkpoint_name = "checkpoints/latest_model_" + args[
"model"] + "_" + basename(normpath(args["dataset"])) + ".ckpt"
if args["continue_training"]:
_log.info('Loaded latest model checkpoint')
saver.restore(sess, model_checkpoint_name)
# Load the data
_log.info("Loading the data ...")
train_input_names, train_output_names, val_input_names, val_output_names, test_input_names, test_output_names = utils.prepare_data(
dataset_dir=args["dataset"])
_log.info("\n***** Begin training *****")
_log.debug("Dataset -->", args["dataset"])
_log.debug("Model -->", args["model"])
_log.debug("Crop Height -->", args["crop_height"])
_log.debug("Crop Width -->", args["crop_width"])
_log.debug("Num Epochs -->", args["num_epochs"])
_log.debug("Batch Size -->", args["batch_size"])
_log.debug("Num Classes -->", num_classes)
_log.debug("Data Augmentation:")
_log.debug("\tVertical Flip -->", args["v_flip"])
_log.debug("\tHorizontal Flip -->", args["h_flip"])
_log.debug("\tBrightness Alteration -->", args["brightness"])
_log.debug("\tRotation -->", args["rotation"])
avg_loss_per_epoch = []
avg_scores_per_epoch = []
avg_iou_per_epoch = []
# Which validation images do we want
val_indices = []
num_vals = min(args["num_val_images"], len(val_input_names))
# Set random seed to make sure models are validated on the same validation images.
# So you can compare the results of different models more intuitively.
random.seed(16)
val_indices = random.sample(range(0, len(val_input_names)), num_vals)
# Do the training here
for epoch in range(args["epoch_start_i"], args["num_epochs"]):
current_losses = []
cnt = 0
# Equivalent to shuffling
id_list = np.random.permutation(len(train_input_names))
num_iters = int(np.floor(len(id_list) / args["batch_size"]))
st = time.time()
epoch_st = time.time()
for i in range(num_iters):
# st=time.time()
input_image_batch = []
output_image_batch = []
# Collect a batch of images
for j in range(args["batch_size"]):
index = i * args["batch_size"] + j
id = id_list[index]
input_image = utils.load_image(train_input_names[id],
args["crop_width"],
args["crop_height"])
output_image = utils.load_image(train_output_names[id],
args["crop_width"],
args["crop_height"])
with tf.device('/cpu:0'):
input_image, output_image = data_augmentation(
input_image, output_image)
# Prep the data. Make sure the labels are in one-hot format
input_image = np.float32(input_image) / 255.0
output_image = np.float32(
helpers.one_hot_it(label=output_image,
label_values=label_values))
input_image_batch.append(
np.expand_dims(input_image, axis=0))
output_image_batch.append(
np.expand_dims(output_image, axis=0))
if args["batch_size"] == 1:
input_image_batch = input_image_batch[0]
output_image_batch = output_image_batch[0]
else:
input_image_batch = np.squeeze(
np.stack(input_image_batch, axis=1))
output_image_batch = np.squeeze(
np.stack(output_image_batch, axis=1))
# Do the training
_, current = sess.run([opt, loss],
feed_dict={
net_input: input_image_batch,
net_output: output_image_batch
})
current_losses.append(current)
cnt = cnt + args["batch_size"]
if cnt % 20 == 0:
string_print = "Epoch = %d Count = %d Iter:(%d of %d) Current_Loss = %.4f Time = %.2f" % (
epoch, cnt, i, num_iters, current, time.time() - st)
utils.LOG(string_print)
st = time.time()
mean_loss = np.mean(current_losses)
avg_loss_per_epoch.append(mean_loss)
# Create directories if needed
if not os.path.isdir("%s/%04d" % ("checkpoints", epoch)):
os.makedirs("%s/%04d" % ("checkpoints", epoch))
# Save latest checkpoint to same file name
_log.info("Saving latest checkpoint")
saver.save(sess, model_checkpoint_name)
if val_indices != 0 and epoch % args["checkpoint_step"] == 0:
_log.info("Saving checkpoint for this epoch")
saver.save(sess, "%s/%04d/model.ckpt" % ("checkpoints", epoch))
if epoch % args["validation_step"] == 0:
_log.info("Performing validation")
target_path = "%s/%04d/val_scores.csv" % ("checkpoints", epoch)
target = open(target_path, 'w')
target.write(
"val_name, avg_accuracy, precision, recall, f1 score, mean iou, %s\n"
% (class_names_string))
scores_list = []
class_scores_list = []
precision_list = []
recall_list = []
f1_list = []
iou_list = []
# Do the validation on a small set of validation images
for ind in val_indices:
input_image = np.expand_dims(np.float32(
utils.load_image(val_input_names[ind], args["crop_width"],
args["crop_height"])
[:args["crop_height"], :args["crop_width"]]),
axis=0) / 255.0
gt = utils.load_image(
val_output_names[ind], args["crop_width"],
args["crop_height"]
)[:args["crop_height"], :args["crop_width"]]
gt = helpers.reverse_one_hot(
helpers.one_hot_it(gt, label_values))
# st = time.time()
output_image = sess.run(network,
feed_dict={net_input: input_image})
output_image = np.array(output_image[0, :, :, :])
output_image = helpers.reverse_one_hot(output_image)
out_vis_image = helpers.colour_code_segmentation(
output_image, label_values)
accuracy, class_accuracies, prec, rec, f1, iou = utils.evaluate_segmentation(
pred=output_image, label=gt, num_classes=num_classes)
file_name = utils.filepath_to_name(val_input_names[ind])
target.write("%s, %f, %f, %f, %f, %f" %
(file_name, accuracy, prec, rec, f1, iou))
for item in class_accuracies:
target.write(", %f" % (item))
target.write("\n")
scores_list.append(accuracy)
class_scores_list.append(class_accuracies)
precision_list.append(prec)
recall_list.append(rec)
f1_list.append(f1)
iou_list.append(iou)
gt = helpers.colour_code_segmentation(gt, label_values)
file_name = os.path.basename(val_input_names[ind])
file_name = os.path.splitext(file_name)[0]
pred_img_path = "%s/%04d/%s_pred.png" % ("checkpoints", epoch,
file_name)
gt_img_path = "%s/%04d/%s_gt.png" % ("checkpoints", epoch,
file_name)
cv2.imwrite(
pred_img_path,
cv2.cvtColor(np.uint8(out_vis_image), cv2.COLOR_RGB2BGR))
cv2.imwrite(gt_img_path,
cv2.cvtColor(np.uint8(gt), cv2.COLOR_RGB2BGR))
# Send the first 16 images to sacred
if ind in val_indices[:16]:
ex.add_artifact(gt_img_path, "GtImage_%d" % ind)
ex.add_artifact(pred_img_path, "PredImage_%d" % ind)
target.close()
ex.add_artifact(target_path)
avg_score = np.mean(scores_list)
class_avg_scores = np.mean(class_scores_list, axis=0)
avg_scores_per_epoch.append(avg_score)
avg_precision = np.mean(precision_list)
avg_recall = np.mean(recall_list)
avg_f1 = np.mean(f1_list)
avg_iou = np.mean(iou_list)
avg_iou_per_epoch.append(avg_iou)
# Sacred info dict gets sent every heartbeat (10s)
ex.info["avg_score"] = avg_score
ex.info["class_avg_scores"] = class_avg_scores
ex.info["avg_precision"] = avg_precision
ex.info["avg_recall"] = avg_recall
ex.info["avg_f1"] = avg_f1
ex.info["avg_iou"] = avg_iou
_log.debug("\nAverage validation accuracy for epoch # %04d = %f" %
(epoch, avg_score))
_log.debug(
"Average per class validation accuracies for epoch # %04d:" %
(epoch))
for index, item in enumerate(class_avg_scores):
_log.debug("%s = %f" % (class_names_list[index], item))
_log.debug("Validation precision = ", avg_precision)
_log.debug("Validation recall = ", avg_recall)
_log.debug("Validation F1 score = ", avg_f1)
_log.debug("Validation IoU score = ", avg_iou)
epoch_time = time.time() - epoch_st
remain_time = epoch_time * (args["num_epochs"] - 1 - epoch)
m, s = divmod(remain_time, 60)
h, m = divmod(m, 60)
if s != 0:
train_time = "Remaining training time = %d hours %d minutes %d seconds\n" % (
h, m, s)
else:
train_time = "Remaining training time : Training completed.\n"
utils.LOG(train_time)
scores_list = []
fig1, ax1 = plt.subplots(figsize=(11, 8))
ax1.plot(range(epoch + 1), avg_scores_per_epoch)
ax1.set_title("Average validation accuracy vs epochs")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Avg. val. accuracy")
plt.savefig('accuracy_vs_epochs.png')
ex.add_artifact("accuracy_vs_epochs.png")
plt.clf()
fig2, ax2 = plt.subplots(figsize=(11, 8))
ax2.plot(range(epoch + 1), avg_loss_per_epoch)
ax2.set_title("Average loss vs epochs")
ax2.set_xlabel("Epoch")
ax2.set_ylabel("Current loss")
plt.savefig('loss_vs_epochs.png')
ex.add_artifact("loss_vs_epochs.png")
plt.clf()
fig3, ax3 = plt.subplots(figsize=(11, 8))
ax3.plot(range(epoch + 1), avg_iou_per_epoch)
ax3.set_title("Average IoU vs epochs")
ax3.set_xlabel("Epoch")
ax3.set_ylabel("Current IoU")
plt.savefig('iou_vs_epochs.png')
ex.add_artifact("iou_vs_epochs.png")
class_names_string = class_names_string + class_name + ", "
else:
class_names_string = class_names_string + class_name
num_classes = len(label_values)
# Initializing the network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess=tf.Session(config=config)
net_input = tf.placeholder(tf.float32,shape=[None,None,None,3])
net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes])
net_edge = tf.placeholder(tf.float32,shape=[None,None])
network, _, output_edge = model_builder.build_model("tbnet", net_input=net_input, num_classes=num_classes, is_training=False)
sess.run(tf.global_variables_initializer())
# Load the weights from the checkpoint
saver=tf.train.Saver(max_to_keep=300)
saver.restore(sess, args.checkpoint_path)
# Load the data
train_input_names,train_output_names, val_input_names, val_output_names, test_input_names, test_output_names, train_edge_names, val_edge_names, test_edge_names = utils.prepare_data(dataset_dir=args.dataset)
# Create directories if needed
if not os.path.isdir("%s"%("Results")):
os.makedirs("%s"%("Results"))
class_scores_list = []
print("Dataset -->", args.dataset)
print("Model -->", args.model)
print("Crop Height -->", args.crop_height)
print("Crop Width -->", args.crop_width)
print("Num Classes -->", num_classes)
print("Image -->", args.image)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess=tf.Session(config=config)
net_input = tf.placeholder(tf.float32,shape=[None,None,None,3])
net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes])
network, _ = model_builder.build_model(model_name=args.model, frontend=args.frontend, net_input=net_input, num_classes=num_classes, crop_width=args.crop_width + 27, crop_height=args.crop_height + 27, is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver=tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
for ind in pre_indices:
print("Testing image " + pre_input_names[ind])
loaded_image = utils.load_image(pre_input_names[ind])
padding_image = cv2.copyMakeBorder(loaded_image, 14, 13, 14, 13, cv2.BORDER_REFLECT)
input_image = np.expand_dims(np.float32(padding_image),axis=0)/255.0
st = time.time()
print("Model -->", args.model)
print("Crop Height -->", args.crop_height)
print("Crop Width -->", args.crop_width)
print("Num Classes -->", num_classes)
print("Image -->", args.image)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes])
network, _ = model_builder.build_model(args.model,
net_input=net_input,
num_classes=num_classes,
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver = tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
print("Testing image " + args.image)
loaded_image = utils.load_image(args.image)
resized_image = cv2.resize(loaded_image, (args.crop_width, args.crop_width))
input_image = np.expand_dims(np.float32(
resized_image[:args.crop_height, :args.crop_width]),
axis=0) / 255.0
class_names_string = class_names_string + class_name
num_classes = len(label_values)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# Compute your softmax cross entropy loss
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes])
network, init_fn = model_builder.build_model(model_name=args.model,
frontend=args.frontend,
net_input=net_input,
num_classes=num_classes,
crop_width=args.crop_width,
crop_height=args.crop_height,
is_training=True)
loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output))
opt = tf.train.RMSPropOptimizer(learning_rate=0.0001, decay=0.995).minimize(
loss, var_list=[var for var in tf.trainable_variables()])
saver = tf.train.Saver(max_to_keep=1000)
sess.run(tf.global_variables_initializer())
utils.count_params()
def train_model(args):
"""
args:
args: global arguments
"""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
print("=====> input size:{}".format(input_size))
print(args)
if args.cuda:
print("=====> use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
# set the seed
setup_seed(GLOBAL_SEED)
print("=====> set Global Seed: ", GLOBAL_SEED)
cudnn.enabled = True
print("=====> building network")
# build the model and initialization
model = build_model(args.model, num_classes=args.classes)
init_weight(model,
nn.init.kaiming_normal_,
nn.BatchNorm2d,
1e-3,
0.1,
mode='fan_in')
print("=====> computing network parameters and FLOPs")
total_parameters = netParams(model)
print("the number of parameters: %d ==> %.2f M" %
(total_parameters, (total_parameters / 1e6)))
# load data and data augmentation
datas, trainLoader, valLoader = build_dataset_train(
args.dataset, args.classes, input_size, args.batch_size,
args.train_type, False, False, args.num_workers)
args.per_iter = len(trainLoader)
args.max_iter = args.max_epochs * args.per_iter
print('=====> Dataset statistics')
print("data['classWeights']: ", datas['classWeights'])
print('mean and std: ', datas['mean'], datas['std'])
# datas['classWeights'] = np.array([4.044603, 2.0614128, 4.2246304, 6.0238333,
# 10.107266, 8.601249, 8.808282], dtype=np.float32)
# datas['mean'] = [0.5, 0.5, 0.5]
# datas['std'] = [0.2, 0.2, 0.2]
# define loss function, respectively
weight = torch.from_numpy(datas['classWeights'])
if args.dataset == 'pollen':
weight = torch.tensor([1., 1.])
if args.dataset == 'camvid':
criteria = CrossEntropyLoss2d(weight=weight,
ignore_label=args.ignore_label)
elif args.dataset == 'camvid' and args.use_label_smoothing:
criteria = CrossEntropyLoss2dLabelSmooth(
weight=weight, ignore_label=args.ignore_label)
elif args.dataset == 'cityscapes' and args.use_ohem:
min_kept = int(args.batch_size // len(args.gpus) * h * w // 16)
criteria = ProbOhemCrossEntropy2d(use_weight=True,
ignore_label=args.ignore_label,
thresh=0.7,
min_kept=min_kept)
elif args.dataset == 'cityscapes' and args.use_label_smoothing:
criteria = CrossEntropyLoss2dLabelSmooth(
weight=weight, ignore_label=args.ignore_label)
elif args.dataset == 'cityscape' and args.use_lovaszsoftmax:
criteria = LovaszSoftmax(ignore_index=args.ignore_label)
elif args.dataset == 'cityscapes' and args.use_focal:
criteria = FocalLoss2d(weight=weight, ignore_index=args.ignore_label)
elif args.dataset == 'seed':
criteria = CrossEntropyLoss2d(weight=weight,
ignore_label=args.ignore_label)
elif args.dataset == 'remote':
criteria = CrossEntropyLoss2d(weight=weight,
ignore_label=args.ignore_label)
elif args.dataset == 'remote' and args.use_ohem:
min_kept = int(args.batch_size // len(args.gpus) * h * w // 16)
criteria = ProbOhemCrossEntropy2d(use_weight=True,
ignore_label=args.ignore_label,
thresh=0.7,
min_kept=min_kept)
elif args.dataset == 'remote' and args.use_label_smoothing:
criteria = CrossEntropyLoss2dLabelSmooth(
weight=weight, ignore_label=args.ignore_label)
elif args.dataset == 'remote' and args.use_lovaszsoftmax:
criteria = LovaszSoftmax(ignore_index=args.ignore_label)
elif args.dataset == 'remote' and args.use_focal:
criteria = FocalLoss2d(weight=weight, ignore_index=args.ignore_label)
else:
criteria = CrossEntropyLoss2d(weight=weight,
ignore_label=args.ignore_label)
if args.cuda:
criteria = criteria.cuda()
if torch.cuda.device_count() > 1:
print("torch.cuda.device_count()=", torch.cuda.device_count())
args.gpu_nums = torch.cuda.device_count()
model = nn.DataParallel(model).cuda() # multi-card data parallel
else:
args.gpu_nums = 1
print("single GPU for training")
model = model.cuda() # 1-card data parallel
args.savedir = (args.savedir + args.dataset + '/' + args.model + 'bs' +
str(args.batch_size) + 'gpu' + str(args.gpu_nums) + "_" +
str(args.train_type) + '/')
if not os.path.exists(args.savedir):
os.makedirs(args.savedir)
start_epoch = 0
# continue training
if args.resume:
if os.path.isfile(args.resume):
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
print("=====> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=====> no checkpoint found at '{}'".format(args.resume))
model.train()
cudnn.benchmark = True
# cudnn.deterministic = True ## my add
logFileLoc = args.savedir + args.logFile
if os.path.isfile(logFileLoc):
logger = open(logFileLoc, 'a')
else:
logger = open(logFileLoc, 'w')
logger.write("Parameters: %s Seed: %s" %
(str(total_parameters), GLOBAL_SEED))
logger.write("\n%s\t\t%s\t%s\t%s" %
('Epoch', 'Loss(Tr)', 'mIOU (val)', 'lr'))
logger.flush()
# define optimization strategy
if args.optim == 'sgd':
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
momentum=0.9,
weight_decay=1e-4)
elif args.optim == 'adam':
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=1e-4)
elif args.optim == 'radam':
optimizer = RAdam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
betas=(0.90, 0.999),
eps=1e-08,
weight_decay=1e-4)
elif args.optim == 'ranger':
optimizer = Ranger(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
betas=(0.95, 0.999),
eps=1e-08,
weight_decay=1e-4)
elif args.optim == 'adamw':
optimizer = AdamW(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=1e-4)
lossTr_list = []
epoches = []
mIOU_val_list = []
print('=====> beginning training')
for epoch in range(start_epoch, args.max_epochs):
# training
lossTr, lr = train(args, trainLoader, model, criteria, optimizer,
epoch)
lossTr_list.append(lossTr)
# validation
if epoch % 2 == 0 or epoch == (args.max_epochs - 1):
epoches.append(epoch)
mIOU_val, per_class_iu = val(args, valLoader, model)
mIOU_val_list.append(mIOU_val)
# record train information
logger.write("\n%d\t\t%.4f\t\t%.4f\t\t%.7f" %
(epoch, lossTr, mIOU_val, lr))
logger.flush()
print("Epoch : " + str(epoch) + ' Details')
print(
"Epoch No.: %d\tTrain Loss = %.4f\t mIOU(val) = %.4f\t lr= %.6f\n"
% (epoch, lossTr, mIOU_val, lr))
else:
# record train information
logger.write("\n%d\t\t%.4f\t\t\t\t%.7f" % (epoch, lossTr, lr))
logger.flush()
print("Epoch : " + str(epoch) + ' Details')
print("Epoch No.: %d\tTrain Loss = %.4f\t lr= %.6f\n" %
(epoch, lossTr, lr))
# save the model
model_file_name = args.savedir + '/model_' + str(epoch + 1) + '.pth'
state = {"epoch": epoch + 1, "model": model.state_dict()}
# Individual Setting for save model !!!
if args.dataset == 'camvid':
torch.save(state, model_file_name)
elif args.dataset == 'cityscapes':
if epoch >= args.max_epochs - 10:
torch.save(state, model_file_name)
elif not epoch % 50:
torch.save(state, model_file_name)
elif args.dataset == 'seed':
torch.save(state, model_file_name)
else:
torch.save(state, model_file_name)
# draw plots for visualization
if epoch % 5 == 0 or epoch == (args.max_epochs - 1):
# Plot the figures per 50 epochs
fig1, ax1 = plt.subplots(figsize=(11, 8))
ax1.plot(range(start_epoch, epoch + 1), lossTr_list)
ax1.set_title("Average training loss vs epochs")
ax1.set_xlabel("Epochs")
ax1.set_ylabel("Current loss")
plt.savefig(args.savedir + "loss_vs_epochs.png")
plt.clf()
fig2, ax2 = plt.subplots(figsize=(11, 8))
ax2.plot(epoches, mIOU_val_list, label="Val IoU")
ax2.set_title("Average IoU vs epochs")
ax2.set_xlabel("Epochs")
ax2.set_ylabel("Current IoU")
plt.legend(loc='lower right')
plt.savefig(args.savedir + "iou_vs_epochs.png")
plt.close('all')
logger.close()
def train_model(args):
"""
args:
args: global arguments
"""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
print("=====> input size:{}".format(input_size))
print(args)
if args.cuda:
print("=====> use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception("No GPU found or Wrong gpu id, please run without --cuda")
# set the seed
setup_seed(GLOBAL_SEED)
print("=====> set Global Seed: ", GLOBAL_SEED)
cudnn.enabled = True
# build the model and initialization
model = build_model(args.model, num_classes=args.classes)
init_weight(model, nn.init.kaiming_normal_,
nn.BatchNorm2d, 1e-3, 0.1,
mode='fan_in')
print("=====> computing network parameters and FLOPs")
total_paramters = netParams(model)
print("the number of parameters: %d ==> %.2f M" % (total_paramters, (total_paramters / 1e6)))
# load data and data augmentation
datas, trainLoader, valLoader = build_dataset_train(args.dataset, input_size, args.batch_size, args.train_type,
args.random_scale, args.random_mirror, args.num_workers)
print('=====> Dataset statistics')
print("data['classWeights']: ", datas['classWeights'])
print('mean and std: ', datas['mean'], datas['std'])
# define loss function, respectively
weight = torch.from_numpy(datas['classWeights'])
if args.dataset == 'camvid':
criteria = CrossEntropyLoss2d(weight=weight, ignore_label=ignore_label)
elif args.dataset == 'cityscapes':
min_kept = int(args.batch_size // len(args.gpus) * h * w // 16)
criteria = ProbOhemCrossEntropy2d(use_weight=True, ignore_label=ignore_label,
thresh=0.7, min_kept=min_kept)
elif args.dataset == 'paris':
criteria = CrossEntropyLoss2d(weight=weight, ignore_label=ignore_label)
# criteria = nn.CrossEntropyLoss(weight=weight)
# min_kept = int(args.batch_size // len(args.gpus) * h * w // 16)
# criteria = ProbOhemCrossEntropy2d(ignore_label=ignore_label, thresh=0.7, min_kept=min_kept, use_weight=False)
elif args.dataset == 'austin':
criteria = BinCrossEntropyLoss2d(weight=weight)
elif args.dataset == 'road':
criteria = BinCrossEntropyLoss2d(weight=weight)
else:
raise NotImplementedError(
"This repository now supports two datasets: cityscapes and camvid, %s is not included" % args.dataset)
if args.cuda:
criteria = criteria.cuda()
if torch.cuda.device_count() > 1:
print("torch.cuda.device_count()=", torch.cuda.device_count())
args.gpu_nums = torch.cuda.device_count()
model = nn.DataParallel(model).cuda() # multi-card data parallel
else:
args.gpu_nums = 1
print("single GPU for training")
model = model.cuda() # 1-card data parallel
args.savedir = (args.savedir + args.dataset + '/' + args.model + 'bs'
+ str(args.batch_size) + 'gpu' + str(args.gpu_nums) + "_" + str(args.train_type) + '/' + str(date) +'/')
if not os.path.exists(args.savedir):
os.makedirs(args.savedir)
start_epoch = 0
# continue training
if args.resume:
if os.path.isfile(args.resume):
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
print("=====> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
else:
print("=====> no checkpoint found at '{}'".format(args.resume))
model.train()
cudnn.benchmark = True
logFileLoc = args.savedir + args.logFile
if os.path.isfile(logFileLoc):
logger = open(logFileLoc, 'a')
else:
logger = open(logFileLoc, 'w')
logger.write("Parameters: %s Seed: %s\n %s\n" % (str(total_paramters/ 1e6), GLOBAL_SEED, args))
logger.write("\n%s\t\t%s\t\t%s\t\t%s\t%s\t%s" % ('Epoch', ' lr', ' Loss', ' Pa', ' Mpa', ' mIOU'))
for i in range(args.classes):
logger.write("\t%s" % ('Class'+str(i)))
logger.flush()
# define optimization criteria
if args.dataset == 'camvid':
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()), args.lr, (0.9, 0.999), eps=1e-08, weight_decay=2e-4)
elif args.dataset == 'cityscapes':
optimizer = torch.optim.SGD(
filter(lambda p: p.requires_grad, model.parameters()), args.lr, momentum=0.9, weight_decay=1e-4)
elif args.dataset == 'paris':
optimizer = torch.optim.SGD(
filter(lambda p: p.requires_grad, model.parameters()), args.lr, momentum=0.9, weight_decay=1e-4)
elif args.dataset == 'austin':
optimizer = torch.optim.SGD(
filter(lambda p: p.requires_grad, model.parameters()), args.lr, momentum=0.9, weight_decay=1e-4)
elif args.dataset == 'road':
optimizer = torch.optim.SGD(
filter(lambda p: p.requires_grad, model.parameters()), args.lr, momentum=0.9, weight_decay=1e-4)
lossTr_list = []
epoches = []
mIOU_val_list = []
max_miou = 0
miou = 0
print('***********************************************\n'
'******* Begining traing *******\n'
'***********************************************')
for epoch in range(start_epoch, args.max_epochs):
# training
lossTr, lr = train(args, trainLoader, model, criteria, optimizer, epoch)
lossTr_list.append(lossTr)
# validation
if (epoch % args.val_epochs == 0 and args.train_val == 'True') or epoch == (args.max_epochs - 1):
epoches.append(epoch)
miou, iou, fmiou, pa, mpa = val(args, valLoader, model)
mIOU_val_list.append(miou)
# record train information
logger.write("\n %d\t\t%.6f\t%.5f\t\t%.4f\t%0.4f\t%0.4f" % (epoch, lr, lossTr, fmiou, pa, miou))
for i in range(len(iou)):
logger.write("\t%0.4f" % (iou[i]))
logger.flush()
print("Epoch %d\tTrain Loss = %.4f\t mIOU(val) = %.4f\t lr= %.5f\n" % (epoch, lossTr, miou, lr))
else:
# record train information
logger.write("\n%d\t%.6f\t\t%.5f" % (epoch, lr, lossTr))
logger.flush()
print("Epoch %d\tTrain Loss = %.4f\t lr= %.6f\n" % (epoch, lossTr, lr))
# save the model
model_file_name = args.savedir + '/model_' + str(epoch) + '.pth'
state = {"epoch": epoch, "model": model.state_dict()}
if max_miou < miou and epoch >= args.max_epochs - 50:
max_miou = miou
torch.save(state, model_file_name)
elif epoch % args.save_epochs == 0:
torch.save(state, model_file_name)
# draw plots for visualization
if epoch % args.val_epochs == 0 or epoch == (args.max_epochs - 1):
# Plot the figures per args.val_epochs epochs
fig1, ax1 = plt.subplots(figsize=(11, 8))
ax1.plot(range(start_epoch, epoch + 1), lossTr_list)
ax1.set_title("Average training loss vs epochs")
ax1.set_xlabel("Epochs")
ax1.set_ylabel("Current loss")
plt.savefig(args.savedir + "loss_vs_epochs.png")
plt.clf()
fig2, ax2 = plt.subplots(figsize=(11, 8))
ax2.plot(epoches, mIOU_val_list, label="Val IoU")
ax2.set_title("Average IoU vs epochs")
ax2.set_xlabel("Epochs")
ax2.set_ylabel("Current IoU")
plt.legend(loc='lower right')
plt.savefig(args.savedir + "iou_vs_epochs.png")
plt.close('all')
logger.close()
def load_model(args, csv_file):
print("Retrieving dataset information ...")
class_names_list, label_values = helpers.get_label_info(os.path.join(args.dataset, "class_dict.csv"))
class_names_string = ""
for class_name in class_names_list:
if not class_name == class_names_list[-1]:
class_names_string = class_names_string + class_name + ", "
else:
class_names_string = class_names_string + class_name
num_classes = len(label_values)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess=tf.Session(config=config)
net_input = tf.placeholder(tf.float32,shape=[None,None,None,3])
net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes])
network, _ = model_builder.build_model(args.model, net_input=net_input, num_classes=num_classes, crop_width=args.crop_width, crop_height=args.crop_height, is_training=False)
#network=tf.nn.softmax(logits=network, axis = 3)
#loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output))
#loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output))
sess.run(tf.global_variables_initializer())
#checkpoint_path = "checkpoints/latest_model_" + args.model + "_" + args.dataset + ".ckpt"
checkpoint_path = "checkpoints/latest_model_" + args.model + "_" + args.dataset + ".ckpt"
# checkpoint_path = "DeepLab_V3_Github_pretrained_Models\deeplabv3_cityscapes_train\latest_model_" + args.model + "_" + args.dataset + ".ckpt"
print('Loading model checkpoint weights ...')
saver=tf.train.Saver(max_to_keep=1000)
saver.restore(sess, checkpoint_path)
#print(network)
#net_input = tf.placeholder(tf.float32,shape=[None,None,3])
layer_A, layer_B = network[0], network[1]
classes = get_classes(csv_file)
ScA = np.zeros((len(classes), len(classes)))
ScB = np.zeros((len(classes), len(classes)))
A_shape = 128*128*256
B_shape = 128*128*256
A_scores = np.zeros((A_shape, len(classes)))
B_scores = np.zeros((B_shape, len(classes)))
idx = 0
for c in classes:
in_dir = c + "_Center/"
files = sorted(os.listdir(in_dir))
for f in files:
print("File: " + f)
input_image = np.expand_dims(np.float32(utils.load_image(in_dir+f)), axis = 0)/255.0
#input_image = input_image.reshape((1,input_image.shape[0],input_image.shape[1],input_image.shape[2]))
#print(input_image.dtype)
#print(input_image)
layer_A_output, layer_B_output = sess.run([layer_A, layer_B], feed_dict={net_input: input_image})
layer_A_output = layer_A_output.reshape((-1,1))
layer_B_output = layer_B_output.reshape((-1,1))
A_scores[:, idx] = layer_A_output[:,0]
B_scores[:, idx] = layer_B_output[:,0]
break
idx += 1
#print(A_scores)
#print(A_scores.shape)
#calculate ScA
for i in range(ScA.shape[0]):
for j in range(ScA.shape[1]):
vecA = A_scores[:, i].reshape((A_shape, 1))
vecB = A_scores[:, j].reshape((A_shape, 1))
ScA[i, j] = np.square(LA.norm(vecA - vecB))
#calculate ScB
for i in range(ScB.shape[0]):
for j in range(ScB.shape[1]):
vecA = B_scores[:, i].reshape((B_shape, 1))
vecB = B_scores[:, j].reshape((B_shape, 1))
ScB[i, j] = np.square(LA.norm(vecA - vecB))
return A_scores, B_scores, ScA, ScB
def get_scores(incoming_dir, A_scores, B_scores, ScA, ScB, csv_file, ScA_B, lmbda):
print("Retrieving dataset information ...")
class_names_list, label_values = helpers.get_label_info(os.path.join(args.dataset, "class_dict.csv"))
class_names_string = ""
for class_name in class_names_list:
if not class_name == class_names_list[-1]:
class_names_string = class_names_string + class_name + ", "
else:
class_names_string = class_names_string + class_name
num_classes = len(label_values)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess=tf.Session(config=config)
net_input = tf.placeholder(tf.float32,shape=[None,None,None,3])
net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes])
network, _ = model_builder.build_model(args.model, net_input=net_input, num_classes=num_classes, crop_width=args.crop_width, crop_height=args.crop_height, is_training=False)
#network=tf.nn.softmax(logits=network, axis = 3)
#loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output))
#loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output))
sess.run(tf.global_variables_initializer())
#checkpoint_path = "checkpoints/latest_model_" + args.model + "_" + args.dataset + ".ckpt"
checkpoint_path = "checkpoints/latest_model_" + args.model + "_" + args.dataset + ".ckpt"
# checkpoint_path = "DeepLab_V3_Github_pretrained_Models\deeplabv3_cityscapes_train\latest_model_" + args.model + "_" + args.dataset + ".ckpt"
print('Loading model checkpoint weights ...')
saver=tf.train.Saver(max_to_keep=1000)
saver.restore(sess, checkpoint_path)
A_shape = 128*128*256
B_shape = 128*128*256
classes = get_classes(csv_file)
#print(network)
#net_input = tf.placeholder(tf.float32,shape=[None,None,3])
layer_A, layer_B, net = network[0], network[1], network[2]
files = sorted(os.listdir(incoming_dir))
print(net)
scores = []
for f in files:
print("File: " + f)
input_image = np.expand_dims(np.float32(utils.load_image(incoming_dir+f)), axis = 0)/255.0
#input_image = input_image.reshape((1,input_image.shape[0],input_image.shape[1],input_image.shape[2]))
#print(input_image.dtype)
#print(input_image)
layer_A_output, layer_B_output, net_output = sess.run([layer_A, layer_B, net], feed_dict={net_input: input_image})
layer_A_output = layer_A_output.reshape((-1,1))
layer_B_output = layer_B_output.reshape((-1,1))
SxA = np.zeros((len(classes), 1))
SxB = np.zeros((len(classes), 1))
idx = 0
for c in range(len(classes)):
vecA = A_scores[:, idx].reshape(A_shape, 1)
vecB = B_scores[:, idx].reshape(B_shape, 1)
SxA[idx, :] = np.square(LA.norm(layer_A_output - vecA))
SxB[idx, :] = np.square(LA.norm(layer_B_output - vecB))
idx += 1
net_output = np.array(net_output[0,:,:,:])
net_output = helpers.reverse_one_hot(net_output)
#print(net_output.shape)
alphas = helpers.get_alpha(net_output)
#print(alphas)
SxA_B = calculate_diff_S(SxA, SxB)
SxA_B_cap = np.zeros((len(classes), 1))
#print(ScA_B)
for i in range(len(classes)):
vecB = ScA_B[:,i].reshape((len(classes), 1))
#print(vecB)
#print(alphas[:,i])
#print(vecB)
SxA_B_cap += alphas[:,i] * vecB
#break
#print(SxA_B_cap)
tau, _ = stats.kendalltau(SxA_B, SxA_B_cap)
#print(tau)
distinctiveness = (1-tau)/2.0
uncertainty = 0.0
for i in range(len(alphas)):
uncertainty += alphas[:,i] * (1-alphas[:,i])
uncertainty = -1.0 * uncertainty
score = (1-lmbda) * distinctiveness + lmbda * uncertainty
print(score[0])
scores.append(score[0])
return scores
print("Crop Width -->", args.crop_width)
print("Num Classes -->", num_classes)
print("Image path -->", args.imagepath)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes])
network, _ = model_builder.build_model(args.model,
net_input=net_input,
num_classes=17,
crop_width=args.crop_width,
crop_height=args.crop_height,
is_training=False,
frontend="MobileNetV2")
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver = tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
for img in os.listdir(args.imagepath):
print("Testing image " + img)
image = os.path.join(args.imagepath, img)
loaded_image = utils.load_image(image)
img_shape = loaded_image.shape
class_names_string = class_names_string + class_name + ", "
else:
class_names_string = class_names_string + class_name
num_classes = len(label_values)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess=tf.Session(config=config)
net_input = tf.placeholder(tf.float32,shape=[None,None,None,3])
net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes])
net_edge = tf.placeholder(tf.float32,shape=[None,None])
network, init_fn, output_edge = model_builder.build_model(model_name="tbnet", frontend="ResNet101", net_input=net_input, num_classes=num_classes, is_training=True)
# The softmax cross entropy loss with the weighting mechanism
print("Computing class weights for", args.dataset, "...")
class_weights = utils.compute_class_weights(labels_dir=args.dataset + "/train_labels", label_values=label_values)
weights = tf.reduce_sum(class_weights * net_output, axis=-1)
unweighted_loss = tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output)
weighted_loss = unweighted_loss * weights
loss_1 = tf.reduce_sum(weighted_loss)
# The binary cross entropy loss
loss_2 = tf.reduce_sum(tf.keras.losses.binary_crossentropy(net_edge, output_edge))
loss = loss_1 + loss_2
opt = tf.train.RMSPropOptimizer(learning_rate=0.00001, decay=0.995).minimize(loss, var_list=[var for var in tf.trainable_variables()])
def train_model(args):
"""
args:
args: global arguments
"""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
print("input size:{}".format(input_size))
print(args)
if args.cuda:
print("use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
# set the seed
setup_seed(GLOBAL_SEED)
print("set Global Seed: ", GLOBAL_SEED)
cudnn.enabled = True
print("building network")
# build the model and initialization
model = build_model(args.model, num_classes=args.classes)
init_weight(model,
nn.init.kaiming_normal_,
nn.BatchNorm2d,
1e-3,
0.1,
mode='fan_in')
print("computing network parameters and FLOPs")
total_paramters = netParams(model)
print("the number of parameters: %d ==> %.2f M" %
(total_paramters, (total_paramters / 1e6)))
# load data and data augmentation
datas, trainLoader, valLoader = build_dataset_train(
args.dataset, input_size, args.batch_size, args.train_type,
args.random_scale, args.random_mirror, args.num_workers)
args.per_iter = len(trainLoader)
args.max_iter = args.max_epochs * args.per_iter
print('Dataset statistics')
print("data['classWeights']: ", datas['classWeights'])
print('mean and std: ', datas['mean'], datas['std'])
# define loss function, respectively
weight = torch.from_numpy(datas['classWeights'])
if args.dataset == 'camvid':
criteria = CrossEntropyLoss2d(weight=weight, ignore_label=ignore_label)
elif args.dataset == 'camvid' and args.use_label_smoothing:
criteria = CrossEntropyLoss2dLabelSmooth(weight=weight,
ignore_label=ignore_label)
elif args.dataset == 'cityscapes' and args.use_ohem:
min_kept = int(args.batch_size // len(args.gpus) * h * w // 16)
criteria = ProbOhemCrossEntropy2d(use_weight=True,
ignore_label=ignore_label,
thresh=0.7,
min_kept=min_kept)
elif args.dataset == 'cityscapes' and args.use_label_smoothing:
criteria = CrossEntropyLoss2dLabelSmooth(weight=weight,
ignore_label=ignore_label)
elif args.dataset == 'cityscapes' and args.use_lovaszsoftmax:
criteria = LovaszSoftmax(ignore_index=ignore_label)
elif args.dataset == 'cityscapes' and args.use_focal:
criteria = FocalLoss2d(weight=weight, ignore_index=ignore_label)
elif args.dataset == 'paris':
criteria = CrossEntropyLoss2d(weight=weight, ignore_label=ignore_label)
else:
raise NotImplementedError(
"This repository now supports two datasets: cityscapes and camvid, %s is not included"
% args.dataset)
if args.cuda:
criteria = criteria.cuda()
if torch.cuda.device_count() > 1:
print("torch.cuda.device_count()=", torch.cuda.device_count())
args.gpu_nums = torch.cuda.device_count()
model = nn.DataParallel(model).cuda() # multi-card data parallel
else:
args.gpu_nums = 1
print("single GPU for training")
model = model.cuda() # 1-card data parallel
args.savedir = (args.savedir + args.dataset + '/' + args.model + 'bs' +
str(args.batch_size) + 'gpu' + str(args.gpu_nums) + "_" +
str(args.train_type) + '/')
if not os.path.exists(args.savedir):
os.makedirs(args.savedir)
with open(args.savedir + 'args.txt', 'w') as f:
f.write('mean:{}\nstd:{}\n'.format(datas['mean'], datas['std']))
f.write("Parameters: {} Seed: {}\n".format(str(total_paramters),
GLOBAL_SEED))
f.write(str(args))
start_epoch = 0
# continue training
if args.resume:
if os.path.isfile(args.resume):
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
print("loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("no checkpoint found at '{}'".format(args.resume))
model.train()
cudnn.benchmark = True
# cudnn.deterministic = True ## my add
# initialize the early_stopping object
early_stopping = EarlyStopping(patience=50)
logFileLoc = args.savedir + args.logFile
if os.path.isfile(logFileLoc):
logger = open(logFileLoc, 'a')
else:
logger = open(logFileLoc, 'w')
logger.write("%s\t%s\t\t%s\t%s\t%s" %
('Epoch', ' lr', 'Loss(Tr)', 'Loss(Val)', 'mIOU(Val)'))
logger.flush()
# define optimization strategy
if args.optim == 'sgd':
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
momentum=0.9,
weight_decay=1e-4)
elif args.optim == 'adam':
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=1e-4)
elif args.optim == 'radam':
optimizer = RAdam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
betas=(0.90, 0.999),
eps=1e-08,
weight_decay=1e-4)
elif args.optim == 'ranger':
optimizer = Ranger(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
betas=(0.95, 0.999),
eps=1e-08,
weight_decay=1e-4)
elif args.optim == 'adamw':
optimizer = AdamW(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=1e-4)
lossTr_list = []
epoches = []
mIOU_val_list = []
lossVal_list = []
print('>>>>>>>>>>>beginning training>>>>>>>>>>>')
for epoch in range(start_epoch, args.max_epochs):
# training
lossTr, lr = train(args, trainLoader, model, criteria, optimizer,
epoch)
lossTr_list.append(lossTr)
# validation
if epoch % args.val_miou_epochs == 0:
epoches.append(epoch)
val_loss, mIOU_val, per_class_iu = val(args, valLoader, criteria,
model, epoch)
mIOU_val_list.append(mIOU_val)
lossVal_list.append(val_loss.item())
# record train information
logger.write(
"\n%d\t%.6f\t%.4f\t\t%.4f\t%0.4f\t %s" %
(epoch, lr, lossTr, val_loss, mIOU_val, str(per_class_iu)))
logger.flush()
print(
"Epoch %d\tlr= %.6f\tTrain Loss = %.4f\tVal Loss = %.4f\tmIOU(val) = %.4f\tper_class_iu= %s\n"
% (epoch, lr, lossTr, val_loss, mIOU_val, str(per_class_iu)))
else:
# record train information
val_loss = val(args, valLoader, criteria, model, epoch)
lossVal_list.append(val_loss.item())
logger.write("\n%d\t%.6f\t%.4f\t\t%.4f" %
(epoch, lr, lossTr, val_loss))
logger.flush()
print("Epoch %d\tlr= %.6f\tTrain Loss = %.4f\tVal Loss = %.4f\n" %
(epoch, lr, lossTr, val_loss))
# save the model
model_file_name = args.savedir + '/model_' + str(epoch) + '.pth'
state = {"epoch": epoch, "model": model.state_dict()}
# Individual Setting for save model
if epoch >= args.max_epochs - 10:
torch.save(state, model_file_name)
elif epoch % 10 == 0:
torch.save(state, model_file_name)
# draw plots for visualization
if os.path.isfile(args.savedir + "loss.png"):
f = open(args.savedir + 'log.txt', 'r')
next(f)
epoch_list = []
lossTr_list = []
lossVal_list = []
for line in f.readlines():
epoch_list.append(line.strip().split()[0])
lossTr_list.append(line.strip().split()[2])
lossVal_list.append(line.strip().split()[3])
assert len(epoch_list) == len(lossTr_list) == len(lossVal_list)
fig1, ax1 = plt.subplots(figsize=(11, 8))
ax1.plot(range(0, epoch + 1), lossTr_list, label='Train_loss')
ax1.plot(range(0, epoch + 1), lossVal_list, label='Val_loss')
ax1.set_title("Average training loss vs epochs")
ax1.set_xlabel("Epochs")
ax1.set_ylabel("Current loss")
ax1.legend()
plt.savefig(args.savedir + "loss.png")
plt.clf()
else:
fig1, ax1 = plt.subplots(figsize=(11, 8))
ax1.plot(range(0, epoch + 1), lossTr_list, label='Train_loss')
ax1.plot(range(0, epoch + 1), lossVal_list, label='Val_loss')
ax1.set_title("Average training loss vs epochs")
ax1.set_xlabel("Epochs")
ax1.set_ylabel("Current loss")
ax1.legend()
plt.savefig(args.savedir + "loss.png")
plt.clf()
fig2, ax2 = plt.subplots(figsize=(11, 8))
ax2.plot(epoches, mIOU_val_list, label="Val IoU")
ax2.set_title("Average IoU vs epochs")
ax2.set_xlabel("Epochs")
ax2.set_ylabel("Current IoU")
ax2.legend()
plt.savefig(args.savedir + "mIou.png")
plt.close('all')
early_stopping.monitor(monitor=val_loss)
if early_stopping.early_stop:
print("Early stopping and Save checkpoint")
if not os.path.exists(model_file_name):
torch.save(state, model_file_name)
break
logger.close()
print('Speed Time: %.2f ms / iter FPS: %.2f' % (speed_time, fps))
return speed_time, fps
if __name__ == '__main__':
parser = ArgumentParser()
parser.add_argument("--size",
type=str,
default="360,480",
help="input size of model")
parser.add_argument('--num-channels', type=int, default=3)
parser.add_argument('--batch-size', type=int, default=1)
parser.add_argument('--classes', type=int, default=19)
parser.add_argument('--iter', type=int, default=100)
parser.add_argument('--model', type=str, default='FFENet')
parser.add_argument("--gpus",
type=str,
default="0",
help="gpu ids (default: 0)")
args = parser.parse_args()
h, w = map(int, args.size.split(','))
model = build_model(args.model,
num_classes=args.classes,
is_training=False)
compute_speed(model, (args.batch_size, args.num_channels, h, w),
int(args.gpus),
iteration=args.iter)
def test_model(args):
"""
main function for testing
param args: global arguments
return: None
"""
print(args)
mIOU_val_max = 0
if args.cuda:
print("use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception("no GPU found or wrong gpu id, please run without --cuda")
# build the model
model = build_model(args.model, num_classes=args.classes)
if args.cuda:
model = model.cuda() # using GPU for inference
cudnn.benchmark = True
if args.save:
if not os.path.exists(args.save_seg_dir):
os.makedirs(args.save_seg_dir)
# load the test set
datas, testLoader = build_dataset_test(args.dataset, args.num_workers)
if not args.best:
if args.checkpoint:
if os.path.isfile(args.checkpoint):
print("loading checkpoint '{}'".format(args.checkpoint))
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
else:
print("no checkpoint found at '{}'".format(args.checkpoint))
raise FileNotFoundError("no checkpoint found at '{}'".format(args.checkpoint))
print("beginning validation")
print("validation set length: ", len(testLoader))
miou, class_iou, fmiou, pa, mpa = test(args, testLoader, model)
# Get the best test result among the last 10 model records.
else:
if args.checkpoint:
if os.path.isfile(args.checkpoint):
dirname, basename = os.path.split(args.checkpoint)
mIOU_val = []
per_class_iu = []
check_num = []
checkpoint_name = glob.glob(dirname+'/*.pth')
for i in checkpoint_name:
name = i.split('/')[-1].split('_')[-1].split('.')[0]
check_num.append(int(name))
check_num.sort()
for i in check_num:
basename = 'model_' + str(i) + '.pth'
resume = os.path.join(dirname, basename)
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model'])
print("beginning test the:" + basename)
print("validation set length: ", len(testLoader))
miou, class_iou, fmiou, pa, mpa = test(args, testLoader, model)
print('Miou Val is ',miou)
mIOU_val.append(miou)
# index = list(range(epoch - 19, epoch + 1))[np.argmax(mIOU_val)]
index = check_num[np.argmax(mIOU_val)]
print("The best mIoU among the models is", index)
mIOU_val_max = np.max(mIOU_val)
else:
print("no checkpoint found at '{}'".format(args.checkpoint))
raise FileNotFoundError("no checkpoint found at '{}'".format(args.checkpoint))
# Save the result
if not args.best:
model_path = os.path.splitext(os.path.basename(args.checkpoint))
args.logFile = 'test_' + model_path[0] + '.txt'
logFileLoc = os.path.join(os.path.dirname(args.checkpoint), args.logFile)
else:
args.logFile = 'test_' + 'best' + str(index) + '.txt'
logFileLoc = os.path.join(os.path.dirname(args.checkpoint), args.logFile)
# Save the result
if os.path.isfile(logFileLoc):
logger = open(logFileLoc, 'a+')
else:
logger = open(logFileLoc, 'w')
logger.write("Max Mean IoU: %.4f" % mIOU_val_max)
logger.flush()
logger.close()
def test_model(args):
"""
main function for testing
param args: global arguments
return: None
"""
print(args)
if args.cuda:
print("=====> use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception("no GPU found or wrong gpu id, please run without --cuda")
# build the model
model = build_model(args.model, num_classes=args.classes)
if args.cuda:
model = model.cuda() # using GPU for inference
cudnn.benchmark = True
if args.save:
if not os.path.exists(args.save_seg_dir):
os.makedirs(args.save_seg_dir)
# load the test set
datas, testLoader = build_dataset_test(args.dataset, args.num_workers)
if not args.best:
if args.checkpoint:
if os.path.isfile(args.checkpoint):
print("=====> loading checkpoint '{}'".format(args.checkpoint))
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
else:
print("=====> no checkpoint found at '{}'".format(args.checkpoint))
raise FileNotFoundError("no checkpoint found at '{}'".format(args.checkpoint))
print("=====> beginning validation")
print("validation set length: ", len(testLoader))
mIOU_val, per_class_iu = test(args, testLoader, model)
print(mIOU_val)
print(per_class_iu)
# Get the best test result among the last 10 model records.
else:
if args.checkpoint:
if os.path.isfile(args.checkpoint):
dirname, basename = os.path.split(args.checkpoint)
epoch = int(os.path.splitext(basename)[0].split('_')[1])
mIOU_val = []
per_class_iu = []
for i in range(epoch - 9, epoch + 1):
basename = 'model_' + str(i) + '.pth'
resume = os.path.join(dirname, basename)
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model'])
print("=====> beginning test the" + basename)
print("validation set length: ", len(testLoader))
mIOU_val_0, per_class_iu_0 = test(args, testLoader, model)
mIOU_val.append(mIOU_val_0)
per_class_iu.append(per_class_iu_0)
index = list(range(epoch - 9, epoch + 1))[np.argmax(mIOU_val)]
print("The best mIoU among the last 10 models is", index)
print(mIOU_val)
per_class_iu = per_class_iu[np.argmax(mIOU_val)]
mIOU_val = np.max(mIOU_val)
print(mIOU_val)
print(per_class_iu)
else:
print("=====> no checkpoint found at '{}'".format(args.checkpoint))
raise FileNotFoundError("no checkpoint found at '{}'".format(args.checkpoint))
# Save the result
if not args.best:
model_path = os.path.splitext(os.path.basename(args.checkpoint))
args.logFile = 'test_' + model_path[0] + '.txt'
logFileLoc = os.path.join(os.path.dirname(args.checkpoint), args.logFile)
else:
args.logFile = 'test_' + 'best' + str(index) + '.txt'
logFileLoc = os.path.join(os.path.dirname(args.checkpoint), args.logFile)
# Save the result
if os.path.isfile(logFileLoc):
logger = open(logFileLoc, 'a')
else:
logger = open(logFileLoc, 'w')
logger.write("Mean IoU: %.4f" % mIOU_val)
logger.write("\nPer class IoU: ")
for i in range(len(per_class_iu)):
logger.write("%.4f\t" % per_class_iu[i])
logger.flush()
logger.close()
def test_model(args):
"""
main function for testing
param args: global arguments
return: None
"""
print(args)
if args.cuda:
print("=====> use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
# print(args.gpus)
# torch.cuda.set_device(0)
if not torch.cuda.is_available():
raise Exception(
"no GPU found or wrong gpu id, please run without --cuda")
# build the model
model = build_model(args.model, num_classes=args.classes)
if args.cuda:
model = model.cuda() # using GPU for inference
cudnn.benchmark = True
if not os.path.exists(args.save_seg_dir):
os.makedirs(args.save_seg_dir)
# load the test set
datas, testLoader = build_dataset_sliding_test(args.dataset,
args.num_workers,
none_gt=True)
if args.checkpoint:
if os.path.isfile(args.checkpoint):
print("=====> loading checkpoint '{}'".format(args.checkpoint))
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model'])
# model.load_state_dict(convert_state_dict(checkpoint['model']))
else:
print("=====> no checkpoint found at '{}'".format(args.checkpoint))
raise FileNotFoundError("no checkpoint found at '{}'".format(
args.checkpoint))
# print("=====> beginning testing")
miou, class_iou, fmiou, pa, mpa = predict_sliding(
model.eval(),
image=testLoader,
tile_size=(args.tile_size, args.tile_size),
classes=args.classes)
print(
'Miou is: {:.4f}\nClass iou is: {}\nFMiou is: {:.4f}\nPa is: {:.4f}\nMpa is: {:.4f}'
.format(miou, class_iou, fmiou, pa, mpa))
logFileLoc = args.save_seg_dir + '/result.txt'
if os.path.isfile(logFileLoc):
logger = open(logFileLoc, 'a')
else:
logger = open(logFileLoc, 'w')
logger.write("%s\t%s\t%s" % (' Pa', ' Mpa', ' mIOU'))
for i in range(args.classes):
logger.write("\t%s" % ('Class' + str(i)))
logger.write("\n%.4f\t%0.4f\t%0.4f" % (fmiou, pa, miou))
for i in range(len(class_iou)):
logger.write("\t%0.4f" % (class_iou[i]))
logger.flush()
logger.flush()
print("Crop Height -->", args.crop_height)
print("Crop Width -->", args.crop_width)
print("Num Classes -->", num_classes)
print("Image -->", args.image)
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
net_input = tf.placeholder(tf.float32, shape=[None, None, None, 3])
net_output = tf.placeholder(tf.float32, shape=[None, None, None, num_classes])
network, _ = model_builder.build_model(args.model,
net_input=net_input,
num_classes=num_classes,
crop_width=args.crop_width,
crop_height=args.crop_height,
is_training=False)
sess.run(tf.global_variables_initializer())
print('Loading model checkpoint weights')
saver = tf.train.Saver(max_to_keep=1000)
saver.restore(sess, args.checkpoint_path)
print("Testing image " + args.image)
loaded_image = utils.load_image(args.image)
resized_image = cv2.resize(loaded_image, (args.crop_width, args.crop_width))
input_image = np.expand_dims(np.float32(
resized_image[:args.crop_height, :args.crop_width]),
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# Compute your softmax cross entropy loss
net_input = tf.placeholder(
tf.float32, shape=[None, input_size['height'], input_size['width'], 3])
net_output = tf.placeholder(
tf.float32,
shape=[None, input_size['height'], input_size['width'], nb_class])
# load the model
network, init_fn = model_builder.build_model(model_name=args.model,
frontend=args.frontend,
net_input=net_input,
num_classes=nb_class,
image_width=input_size['width'],
image_height=input_size['height'],
is_training=True)
loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(logits=network, labels=net_output))
optimizer = tf.train.RMSPropOptimizer(
learning_rate=args.learning_rate, decay=args.regularization).minimize(
loss, var_list=[var for var in tf.trainable_variables()])
saver = tf.train.Saver(max_to_keep=1000)
sess.run(tf.global_variables_initializer())
# If pre-trained ResNet required, load weights (must be done AFTER variables are initialized with sess.run(tf.global_variables_initializer())
fps = iteration / elapsed_time
print('Elapsed Time: [%.2f s / %d iter]' % (elapsed_time, iteration))
print('Speed Time: %.2f ms / iter FPS: %.2f' % (speed_time, fps))
return speed_time, fps
if __name__ == '__main__':
parser = ArgumentParser()
parser.add_argument("--size",
type=str,
default="512,1024",
help="input size of model")
parser.add_argument('--num-channels', type=int, default=3)
parser.add_argument('--batch-size', type=int, default=1)
parser.add_argument('--classes', type=int, default=19)
parser.add_argument('--iter', type=int, default=100)
parser.add_argument('--model', type=str, default='ENet')
parser.add_argument("--gpus",
type=str,
default="0",
help="gpu ids (default: 0)")
args = parser.parse_args()
h, w = map(int, args.size.split(','))
model = build_model(args.model, num_classes=args.classes)
compute_speed(model, (args.batch_size, args.num_channels, h, w),
int(args.gpus),
iteration=args.iter)
def main(args):
"""
args:
args: global arguments
"""
# set the seed
setup_seed(GLOBAL_SEED)
# cudnn.enabled = True
# cudnn.benchmark = True # find the optimal configuration
# cudnn.deterministic = True # reduce volatility
# learning scheduling, for 10 epoch lr*0.8
# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.85)
# build the model and initialization weights
model = build_model(args.model, args.classes, args.backbone, args.pretrained, args.out_stride, args.mult_grid)
# define loss function, respectively
criterion = build_loss(args, None, ignore_label)
# load train set and data augmentation
datas, traindataset = build_dataset_train(args.root, args.dataset, args.base_size, args.crop_size)
# load the test set, if want set cityscapes test dataset change none_gt=False
testdataset, class_dict_df = build_dataset_test(args.root, args.dataset, args.crop_size,
mode=args.predict_mode, gt=True)
# move model and criterion on cuda
if args.cuda:
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus_id
dist.init_process_group(backend="nccl", init_method='env://')
args.local_rank = torch.distributed.get_rank()
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
gpus = len(list(os.environ["CUDA_VISIBLE_DEVICES"])) - (len(list(os.environ["CUDA_VISIBLE_DEVICES"])) // 2)
trainLoader, model, criterion = Distribute(args, traindataset, model, criterion, device, gpus)
# test with distributed
# testLoader, _, _ = Distribute(args, testdataset, model, criterion, device, gpus)
# test with single card
testLoader = data.DataLoader(testdataset, batch_size=args.batch_size,
shuffle=True, num_workers=args.batch_size, pin_memory=True, drop_last=False)
if not torch.cuda.is_available():
raise Exception("No GPU found or Wrong gpu id, please run without --cuda")
# define optimization strategy
# parameters = [{'params': model.get_1x_lr_params(), 'lr': args.lr},
# {'params': model.get_10x_lr_params(), 'lr': args.lr}]
parameters = model.parameters()
if args.optim == 'sgd':
optimizer = torch.optim.SGD(parameters, lr=args.lr, momentum=0.9, weight_decay=5e-4, nesterov=False)
elif args.optim == 'adam':
optimizer = torch.optim.Adam(parameters, weight_decay=5e-4)
elif args.optim == 'adamw':
optimizer = torch.optim.AdamW(parameters, weight_decay=5e-4)
# initial log file val output save
args.savedir = (args.savedir + args.dataset + '/' + args.model + '/')
if not os.path.exists(args.savedir) and args.local_rank == 0:
os.makedirs(args.savedir)
# save_seg_dir
args.save_seg_dir = os.path.join(args.savedir, args.predict_mode)
if not os.path.exists(args.save_seg_dir) and args.local_rank == 0:
os.makedirs(args.save_seg_dir)
recorder = record_log(args)
if args.resume == None and args.local_rank == 0:
recorder.record_args(datas, str(netParams(model) / 1e6) + ' M', GLOBAL_SEED)
# initialize the early_stopping object
early_stopping = EarlyStopping(patience=300)
start_epoch = 1
if args.local_rank == 0:
print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n"
">>>>>>>>>>> beginning training >>>>>>>>>>>\n"
">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
epoch_list = []
lossTr_list = []
Miou_list = []
lossVal_list = []
Miou = 0
Best_Miou = 0
# continue training
if args.resume:
logger, lines = recorder.resume_logfile()
for index, line in enumerate(lines):
lossTr_list.append(float(line.strip().split()[2]))
if len(line.strip().split()) != 3:
epoch_list.append(int(line.strip().split()[0]))
lossVal_list.append(float(line.strip().split()[3]))
Miou_list.append(float(line.strip().split()[5]))
if os.path.isfile(args.resume):
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch'] + 1
optimizer.load_state_dict(checkpoint['optimizer'])
check_list = [i for i in checkpoint['model'].items()]
# Read weights with multiple cards, and continue training with a single card this time
if 'module.' in check_list[0][0]:
new_stat_dict = {}
for k, v in checkpoint['model'].items():
new_stat_dict[k[:]] = v
model.load_state_dict(new_stat_dict, strict=True)
# Read the training weight of a single card, and continue training with a single card this time
else:
model.load_state_dict(checkpoint['model'])
if args.local_rank == 0:
print("loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
else:
if args.local_rank == 0:
print("no checkpoint found at '{}'".format(args.resume))
else:
logger = recorder.initial_logfile()
logger.flush()
for epoch in range(start_epoch, args.max_epochs + 1):
start_time = time.time()
# training
train_start = time.time()
lossTr, lr = train(args, trainLoader, model, criterion, optimizer, epoch, device)
if args.local_rank == 0:
lossTr_list.append(lossTr)
train_end = time.time()
train_per_epoch_seconds = train_end - train_start
validation_per_epoch_seconds = 60 # init validation time
# validation if mode==validation, predict with label; elif mode==predict, predict without label.
if epoch % args.val_epochs == 0 or epoch == 1 or args.max_epochs - 10 < epoch <= args.max_epochs:
validation_start = time.time()
loss, FWIoU, Miou, MIoU, PerCiou_set, Pa, PerCpa_set, Mpa, MF, F_set, F1_avg = \
predict_multiscale_sliding(args=args, model=model,
testLoader=testLoader,
class_dict_df=class_dict_df,
# scales=[1.25, 1.5, 1.75, 2.0],
scales=[1.0],
overlap=0.3,
criterion=criterion,
mode=args.predict_type,
save_result=True)
torch.cuda.empty_cache()
if args.local_rank == 0:
epoch_list.append(epoch)
Miou_list.append(Miou)
lossVal_list.append(loss.item())
# record trainVal information
recorder.record_trainVal_log(logger, epoch, lr, lossTr, loss,
FWIoU, Miou, MIoU, PerCiou_set, Pa, Mpa,
PerCpa_set, MF, F_set, F1_avg,
class_dict_df)
torch.cuda.empty_cache()
validation_end = time.time()
validation_per_epoch_seconds = validation_end - validation_start
else:
if args.local_rank == 0:
# record train information
recorder.record_train_log(logger, epoch, lr, lossTr)
# # Update lr_scheduler. In pytorch 1.1.0 and later, should call 'optimizer.step()' before 'lr_scheduler.step()'
# lr_scheduler.step()
if args.local_rank == 0:
# draw log fig
draw_log(args, epoch, epoch_list, lossTr_list, Miou_list, lossVal_list)
# save the model
model_file_name = args.savedir + '/best_model.pth'
last_model_file_name = args.savedir + '/last_model.pth'
state = {
"epoch": epoch,
"model": model.state_dict(),
'optimizer': optimizer.state_dict()
}
if Miou > Best_Miou:
Best_Miou = Miou
torch.save(state, model_file_name)
recorder.record_best_epoch(epoch, Best_Miou, Pa)
# early_stopping monitor
early_stopping.monitor(monitor=Miou)
if early_stopping.early_stop:
print("Early stopping and Save checkpoint")
if not os.path.exists(last_model_file_name):
torch.save(state, last_model_file_name)
torch.cuda.empty_cache() # empty_cache
loss, FWIoU, Miou, Miou_Noback, PerCiou_set, Pa, PerCpa_set, Mpa, MF, F_set, F1_Noback = \
predict_multiscale_sliding(args=args, model=model,
testLoader=testLoader,
scales=[1.0],
overlap=0.3,
criterion=criterion,
mode=args.predict_type,
save_result=False)
print("Epoch {} lr= {:.6f} Train Loss={:.4f} Val Loss={:.4f} Miou={:.4f} PerCiou_set={}\n"
.format(epoch, lr, lossTr, loss, Miou, str(PerCiou_set)))
break
total_second = start_time + (args.max_epochs - epoch) * train_per_epoch_seconds + \
((args.max_epochs - epoch) / args.val_epochs + 10) * validation_per_epoch_seconds + 43200
print('Best Validation MIoU:{}'.format(Best_Miou))
print('Training deadline is: {}\n'.format(time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(total_second))))