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))))