def main(): nclass = 7 batch_size = 4 drop_last = True train_loader, val_loader = make_dataloader(batch_size=batch_size, drop_last=drop_last) model = DeepLab(num_classes=nclass, output_stride=16, freeze_bn=False) model = load_pretrained_model(model) evaluator = Evaluator(nclass) model = model.cuda() model.eval() evaluator.reset() for i, sample in enumerate(val_loader): image, target = sample['image'], sample['label'] image, target = image.cuda(), target.cuda() with torch.no_grad(): output = model(image) pred = output.data.cpu().numpy() target = target.cpu().numpy() pred = np.argmax(pred, axis=1) evaluator.add_batch(target, pred) Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format( Acc, Acc_class, mIoU, FWIoU))
def evaluate_single(gt,pred,num_of_class): evaluator = Evaluator(num_of_class) evaluator.reset() evaluator.add_batch(gt,pred) Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() return format(Acc, '.4f'), format(Acc_class, '.4f'), format(mIoU, '.4f'), format(FWIoU, '.4f')
def evaluate_batch(gt_list,pred_list,num_of_class): evaluator = Evaluator(num_of_class) evaluator.reset() for i in range(len(gt_list)): evaluator.add_batch(gt_list[i],pred_list[i]) Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() return format(Acc, '.4f'), format(Acc_class, '.4f'), format(mIoU, '.4f'), format(FWIoU, '.4f')
def evaluate(image_file_paths, gt_file_paths, model, nn_input_size, num_classes, window_sizes=None, step_sizes=None, device=None): """ Evaluate model performance :param image_file_paths: paths to images that will be predicted :param gt_file_paths: paths to ground truth masks :param model: model used for prediction :param nn_input_size: size that the images will be scaled to before feeding them into the nn :param num_classes: number of classes :param window_sizes: list of sizes that determine how the image will be cut (for sliding window). Image will be cut into squares :param step_sizes: list of step sizes for sliding window :param device: PyTorch device (cpu or gpu) :return: list of prediction masks if res_fcn is None, else Nothing """ model.eval() evaluator = Evaluator(num_classes) def ev(prediction, index): pred_class_id_mask = np.argmax(prediction, axis=-1) gt_colour_mask = cv2.imread(gt_file_paths[index])[:, :, 0] gt_class_id_mask = colour_mask_to_class_id_mask(gt_colour_mask) # Add sample into evaluator evaluator.add_batch(gt_class_id_mask, pred_class_id_mask) with torch.no_grad(): predict.predict(image_file_paths=image_file_paths, model=model, nn_input_size=nn_input_size, res_fcn=ev, window_sizes=window_sizes, step_sizes=step_sizes, device=device) with np.errstate(divide='ignore', invalid='ignore'): acc = evaluator.Pixel_Accuracy() acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() fWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() return acc, acc_class, mIoU, fWIoU
def main(): os.environ["CUDA_VISIBLE_DEVICES"] = "3" testdataset = Cityscapesloader('Cityscapes/', split='val', ) testdataloader = DataLoader(testdataset, batch_size=4, shuffle=True, num_workers=8, pin_memory=True) ICnet_ = ICnet_model(19, (512, 1024)).cuda() model_path = 'checkpoints/last.pth' if os.path.exists(model_path): ICnet_.load(model_path) summary(ICnet_, ( 3, 512, 1024)) ICnet_.eval() eval = Evaluator(19) eval.reset() with torch.no_grad(): for batch_index, (image_batch, labels_batch) in enumerate(testdataloader): print(batch_index) image = image_batch.cuda() labels_batch = np.array(labels_batch) output = ICnet_.forward(image).cpu().numpy() pred = np.argmax(output, axis=1) eval.add_batch(labels_batch, pred) Pixel_Accuracy = eval.Pixel_Accuracy() print('Pixel_Accuracy = ', Pixel_Accuracy)
def main(): # 配置设备 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 数据路径 train_dir = './data_road/training' # 超参数 num_epochs = 60 learning_rate = 0.001 img_size = (640, 192) num_class = 2 SAVE_INTERVAL = 5 evaluator = Evaluator(num_class=num_class) # 构建Dataset实例 train_data = LaneDataset(data_dir=train_dir, img_size=img_size) train_loader = DataLoader(train_data, batch_size=2, shuffle=True) # 构建Model实例 model = FCN8s(n_class=2) vgg16_path = "models/vgg16_from_caffe.pth" vgg16 = load_vgg16(model_file=vgg16_path) model.copy_params_from_vgg16(vgg16) model = model.to(device) print("模型加载成功!!!") # 定义损失函数和优化器 criterion = nn.BCELoss() optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4) Acc_list = [] Acc_class_list = [] mIoU_list = [] FWIoU_list = [] loss_list = [] # 训练 for epoch in range(num_epochs): train_loss = 0.0 evaluator.reset() for i, (image, label) in enumerate(train_loader): image = image.to(device) label = label.to(device) optimizer.zero_grad() # 前向传播 output = model(image) output = torch.sigmoid(output) loss = criterion(output, label) output = torch.argmax(output, dim=1).cpu().numpy() label = torch.argmax(label, dim=1).cpu().numpy() # 后向传播及优化 loss.backward() train_loss += loss.item() optimizer.step() evaluator.add_batch(label, output) # 添加output和label,用于后续评估 # 计算像素准确率、平均像素准确率、平均交并比、频权交并比 Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() Acc_list.append(Acc) Acc_class_list.append(Acc_class) mIoU_list.append(mIoU) FWIoU_list.append(FWIoU) loss_list.append(train_loss/len(train_loader)) # 保存模型 if epoch % SAVE_INTERVAL == 0 or (epoch+1) == num_epochs: torch.save(model.state_dict(), './models/fcn8s_{}.pth'.format(epoch)) print("Epoch_{}: train_loss: {:.6f} Acc: {:.4f} Acc_class: {:.4f} mIoU: {:.4f} FWIoU: {:.4f}" .format(epoch+1, train_loss/len(train_loader), Acc, Acc_class, mIoU, FWIoU)) # 绘制曲线 draw_figures(loss_list, title='train_loss') draw_figures(Acc_list, title='Acc') draw_figures(Acc_class_list, title='Acc_class') draw_figures(mIoU_list, title='mIoU') draw_figures(FWIoU_list, title='FWIoU') print("完成曲线绘制!!!")
class Trainer(object): def __init__(self, args): self.args = args # define Dataloader self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(args) # define network model = DeepLab(num_classes=self.nclass, output_stride=args.out_stride) optimizer = torch.optim.Adam(model.parameters(), lr=args.base_lr, weight_decay=args.weight_decay) self.criterion = SegmentationLoss(weight=None, cuda=args.cuda).build_loss(mode=args.loss_type) self.model, self.optimizer = model, optimizer self.evaluator = Evaluator(self.nclass) self.best_pred = 0.0 self.trainloss_history = [] self.valloss_history = [] self.train_plot = [] self.val_plot = [] # every 10 epochs the lr will multiply 0.1 self.scheduler = LR_Scheduler(args.lr_scheduler, args.base_lr, args.epochs, len(self.train_loader), lr_step=20) if args.cuda: self.model = self.model.cuda() def training(self, epoch): train_loss = 0.0 self.model.train() #tbar = tqdm(self.train_loader) # Instantly make your loops show a smart progress meter for i, sample in enumerate(self.train_loader): if i >= 10000: break else: image, target = sample['image'], sample['label'] if self.args.cuda: image, target = image.cuda(), target.cuda() self.scheduler(self.optimizer, i, epoch, self.best_pred) self.optimizer.zero_grad() output = self.model(image) loss = self.criterion(output, target) loss.backward() self.optimizer.step() self.trainloss_history.append(loss.data.cpu().numpy()) if epoch == 3 and i == 1: print("after 3 epochs, the result of training data:") visualize_image("coco", image, target, output) if epoch == 9 and i == 1: print("after 10 epochs, the result of training data:") visualize_image("coco", image, target, output) if epoch == 29 and i == 1: print("after 20 epochs, the result of training data:") visualize_image("coco", image, target, output) if epoch == 29 and i == 3: visualize_image("coco", image, target, output) if epoch == 29 and i == 5: visualize_image("coco", image, target, output) if epoch == 29 and i == 7: visualize_image("coco", image, target, output) if epoch == 29 and i == 9: visualize_image("coco", image, target, output) if epoch == 49 and i == 1: print("after 30 epochs, the result of training data:") visualize_image("coco", image, target, output) if epoch == 49 and i == 3: visualize_image("coco", image, target, output) if epoch == 49 and i == 5: visualize_image("coco", image, target, output) if epoch == 49 and i == 7: visualize_image("coco", image, target, output) if epoch == 49 and i == 9: visualize_image("coco", image, target, output) pred = output.data.cpu().numpy() pred = np.argmax(pred, axis=1) self.evaluator.add_batch(target.cpu().numpy(), pred) acc = self.evaluator.Pixel_Accuracy() #acc_class = self.evaluator.Pixel_Accuracy_Class() mIoU = self.evaluator.Mean_Intersection_over_Union() last_loss = self.trainloss_history[-i:] train_loss = np.mean(last_loss) self.train_plot.append(train_loss) print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size )) print('Loss: %.3f' % train_loss) print("Train_Acc:{}".format(acc), "Train_mIoU:{}".format(mIoU)) if epoch == 20: torch.save(self.model.state_dict(), '/home/wan/Segmentation/model_20') if epoch == 30: torch.save(self.model.state_dict(), '/home/wan/Segmentation/model_30') if epoch == 40: torch.save(self.model.state_dict(), '/home/wan/Segmentation/model_40') if epoch == 50: torch.save(self.model.state_dict(), '/home/wan/Segmentation/model_50') def validation(self, epoch): self.model.eval() self.evaluator.reset() #tbar = tqdm(self.val_loader, desc='\r') val_loss = 0.0 for i, sample in enumerate(self.val_loader): if i >= 750: break else: image, target = sample['image'], sample['label'] if self.args.cuda: image, target = image.cuda(), target.cuda() with torch.no_grad(): output = self.model(image) loss = self.criterion(output, target) self.valloss_history.append(loss.data.cpu().numpy()) pred = output.data.cpu().numpy() pred = np.argmax(pred, axis=1) # visualize the prediction and target if epoch == 3 and i == 0: print("after 3 epochs, the result of training data:") visualize_image("coco", image, target, output) if epoch ==9 and i == 0: print("after 10 epochs, the result of training data:") visualize_image("coco", image, target, output) if epoch == 29 and i == 0: print("after 20 epochs, the result of training data:") visualize_image("coco", image, target, output) if epoch == 29 and i == 1: visualize_image("coco", image, target, output) if epoch == 29 and i == 2: visualize_image("coco", image, target, output) if epoch == 29 and i == 3: visualize_image("coco", image, target, output) if epoch == 29 and i == 4: visualize_image("coco", image, target, output) if epoch == 49 and i == 0: print("after 30 epochs, the result of training data:") visualize_image("coco", image, target, output) if epoch == 49 and i == 1: visualize_image("coco", image, target, output) if epoch == 49 and i == 2: visualize_image("coco", image, target, output) if epoch == 49 and i == 3: visualize_image("coco", image, target, output) if epoch == 49 and i == 4: visualize_image("coco", image, target, output) self.evaluator.add_batch(target.cpu().numpy(), pred) acc = self.evaluator.Pixel_Accuracy() #acc_class = self.evaluator.Pixel_Accuracy_Class() mIoU = self.evaluator.Mean_Intersection_over_Union() last_loss = self.valloss_history[-i:] val_loss = np.mean(last_loss) self.val_plot.append(val_loss) print('Validation: ') print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size )) print("Acc:{}".format(acc), "mIoU:{}".format(mIoU)) print('Loss: %.3f' % val_loss)
class Trainer: def __init__(self, data_train, data_valid, image_base_dir, instructions): """ :param data_train: :param data_valid: :param image_base_dir: :param instructions: """ self.image_base_dir = image_base_dir self.data_valid = data_valid self.instructions = instructions # specify model save dir self.model_name = instructions[STR.MODEL_NAME] # now = time.localtime() # start_time = "{}-{}-{}T{}:{}:{}".format(now.tm_year, now.tm_mon, now.tm_mday, now.tm_hour, now.tm_min, # now.tm_sec) experiment_folder_path = os.path.join(paths.MODELS_FOLDER_PATH, self.model_name) if os.path.exists(experiment_folder_path): Warning( "Experiment folder exists already. Files might be overwritten") os.makedirs(experiment_folder_path, exist_ok=True) # define saver and save instructions self.saver = Saver(folder_path=experiment_folder_path, instructions=instructions) self.saver.save_instructions() # define Tensorboard Summary self.writer = SummaryWriter(log_dir=experiment_folder_path) nn_input_size = instructions[STR.NN_INPUT_SIZE] state_dict_file_path = instructions.get(STR.STATE_DICT_FILE_PATH, None) self.colour_mapping = mapping.get_colour_mapping() # define transformers for training crops_per_image = instructions.get(STR.CROPS_PER_IMAGE, 10) apply_random_cropping = (STR.CROPS_PER_IMAGE in instructions.keys()) and \ (STR.IMAGES_PER_BATCH in instructions.keys()) print("{}applying random cropping".format( "" if apply_random_cropping else "_NOT_ ")) t = [Normalize()] if apply_random_cropping: t.append( RandomCrop(min_size=instructions.get(STR.CROP_SIZE_MIN, 400), max_size=instructions.get(STR.CROP_SIZE_MAX, 1000), crop_count=crops_per_image)) t += [ Resize(nn_input_size), Flip(p_vertical=0.2, p_horizontal=0.5), ToTensor() ] transformations_train = transforms.Compose(t) # define transformers for validation transformations_valid = transforms.Compose( [Normalize(), Resize(nn_input_size), ToTensor()]) # set up data loaders dataset_train = DictArrayDataSet(image_base_dir=image_base_dir, data=data_train, num_classes=len( self.colour_mapping.keys()), transformation=transformations_train) # define batch sizes self.batch_size = instructions[STR.BATCH_SIZE] if apply_random_cropping: self.data_loader_train = DataLoader( dataset=dataset_train, batch_size=instructions[STR.IMAGES_PER_BATCH], shuffle=True, collate_fn=custom_collate) else: self.data_loader_train = DataLoader(dataset=dataset_train, batch_size=self.batch_size, shuffle=True, collate_fn=custom_collate) dataset_valid = DictArrayDataSet(image_base_dir=image_base_dir, data=data_valid, num_classes=len( self.colour_mapping.keys()), transformation=transformations_valid) self.data_loader_valid = DataLoader(dataset=dataset_valid, batch_size=self.batch_size, shuffle=False, collate_fn=custom_collate) self.num_classes = dataset_train.num_classes() # define model print("Building model") self.model = DeepLab(num_classes=self.num_classes, backbone=instructions.get(STR.BACKBONE, "resnet"), output_stride=instructions.get( STR.DEEPLAB_OUTPUT_STRIDE, 16)) # load weights if state_dict_file_path is not None: print("loading state_dict from:") print(state_dict_file_path) load_state_dict(self.model, state_dict_file_path) learning_rate = instructions.get(STR.LEARNING_RATE, 1e-5) train_params = [{ 'params': self.model.get_1x_lr_params(), 'lr': learning_rate }, { 'params': self.model.get_10x_lr_params(), 'lr': learning_rate }] # choose gpu or cpu self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") if instructions.get(STR.MULTI_GPU, False): if torch.cuda.device_count() > 1: print("Using ", torch.cuda.device_count(), " GPUs!") self.model = nn.DataParallel(self.model) self.model.to(self.device) # Define Optimizer self.optimizer = torch.optim.SGD(train_params, momentum=0.9, weight_decay=5e-4, nesterov=False) # calculate class weights if instructions.get(STR.CLASS_STATS_FILE_PATH, None): class_weights = calculate_class_weights( instructions[STR.CLASS_STATS_FILE_PATH], self.colour_mapping, modifier=instructions.get(STR.LOSS_WEIGHT_MODIFIER, 1.01)) class_weights = torch.from_numpy(class_weights.astype(np.float32)) else: class_weights = None self.criterion = SegmentationLosses( weight=class_weights, cuda=self.device.type != "cpu").build_loss() # Define Evaluator self.evaluator = Evaluator(self.num_classes) # Define lr scheduler self.scheduler = None if instructions.get(STR.USE_LR_SCHEDULER, True): self.scheduler = LR_Scheduler(mode="cos", base_lr=learning_rate, num_epochs=instructions[STR.EPOCHS], iters_per_epoch=len( self.data_loader_train)) # print information before training start print("-" * 60) print("instructions") pprint(instructions) model_parameters = sum([p.nelement() for p in self.model.parameters()]) print("Model parameters: {:.2E}".format(model_parameters)) self.best_prediction = 0.0 def train(self, epoch): self.model.train() train_loss = 0.0 # create a progress bar pbar = tqdm(self.data_loader_train) num_batches_train = len(self.data_loader_train) # go through each item in the training data for i, sample in enumerate(pbar): # set input and target nn_input = sample[STR.NN_INPUT].to(self.device) nn_target = sample[STR.NN_TARGET].to(self.device, dtype=torch.float) if self.scheduler: self.scheduler(self.optimizer, i, epoch, self.best_prediction) # run model output = self.model(nn_input) # calc losses loss = self.criterion(output, nn_target) # # save step losses # combined_loss_steps.append(float(loss)) # regression_loss_steps.append(float(regression_loss)) # classification_loss_steps.append(float(classification_loss)) train_loss += loss.item() pbar.set_description('Train loss: %.3f' % (train_loss / (i + 1))) self.writer.add_scalar('train/total_loss_iter', loss.item(), i + num_batches_train * epoch) # calculate gradient and update model weights loss.backward() # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) self.optimizer.step() self.optimizer.zero_grad() self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch) print("[Epoch: {}, num images/crops: {}]".format( epoch, num_batches_train * self.batch_size)) print("Loss: {:.2f}".format(train_loss)) def validation(self, epoch): self.model.eval() self.evaluator.reset() test_loss = 0.0 pbar = tqdm(self.data_loader_valid, desc='\r') num_batches_val = len(self.data_loader_valid) for i, sample in enumerate(pbar): # set input and target nn_input = sample[STR.NN_INPUT].to(self.device) nn_target = sample[STR.NN_TARGET].to(self.device, dtype=torch.float) with torch.no_grad(): output = self.model(nn_input) loss = self.criterion(output, nn_target) test_loss += loss.item() pbar.set_description('Test loss: %.3f' % (test_loss / (i + 1))) pred = output.data.cpu().numpy() pred = np.argmax(pred, axis=1) nn_target = nn_target.cpu().numpy() # Add batch sample into evaluator self.evaluator.add_batch(nn_target, pred) # Fast test during the training Acc = self.evaluator.Pixel_Accuracy() Acc_class = self.evaluator.Pixel_Accuracy_Class() mIoU = self.evaluator.Mean_Intersection_over_Union() FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union() self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch) self.writer.add_scalar('val/mIoU', mIoU, epoch) self.writer.add_scalar('val/Acc', Acc, epoch) self.writer.add_scalar('val/Acc_class', Acc_class, epoch) self.writer.add_scalar('val/fwIoU', FWIoU, epoch) print('Validation:') print("[Epoch: {}, num crops: {}]".format( epoch, num_batches_val * self.batch_size)) print( "Acc:{:.2f}, Acc_class:{:.2f}, mIoU:{:.2f}, fwIoU: {:.2f}".format( Acc, Acc_class, mIoU, FWIoU)) print("Loss: {:.2f}".format(test_loss)) new_pred = mIoU is_best = new_pred > self.best_prediction if is_best: self.best_prediction = new_pred self.saver.save_checkpoint(self.model, is_best, epoch)
def main(): nclass = 7 lr = 0.01 num_epochs = 400 batch_size = 4 best_pred = 0.0 drop_last = True train_loader, val_loader = make_dataloader(batch_size=batch_size, drop_last=drop_last) model = DeepLab(num_classes=nclass, output_stride=16, freeze_bn=False) train_params = [{ 'params': model.get_1x_lr_params(), 'lr': lr }, { 'params': model.get_10x_lr_params(), 'lr': lr * 10 }] optimizer = torch.optim.SGD(train_params, momentum=0.9, weight_decay=5e-4, nesterov=False) criterion = SegmentationLosses(weight=None, cuda=True).build_loss(mode='ce') evaluator = Evaluator(nclass) scheduler = LR_Scheduler(mode='poly', base_lr=lr, num_epochs=num_epochs, iters_per_epoch=len(train_loader)) model = model.cuda() for epoch in range(num_epochs): train_loss = 0.0 count = 0 model.train() for i, sample in enumerate(train_loader): image, target = sample['image'], sample['label'] image, target = image.cuda(), target.cuda() scheduler(optimizer, i, epoch, best_pred) optimizer.zero_grad() output = model(image) loss = criterion(output, target) loss.backward() optimizer.step() train_loss += loss.item() count += 1 train_loss = train_loss / count print( 'Training [Epoch: %d, best_pred: %.4f, numImages: %5d, Loss: %.3f]' % (epoch, best_pred, i * batch_size + image.data.shape[0], train_loss)) writer.add_scalar('scalar/loss_seg_train', train_loss, epoch) model.eval() evaluator.reset() test_loss = 0.0 count = 0 for i, sample in enumerate(val_loader): image, target = sample['image'], sample['label'] image, target = image.cuda(), target.cuda() with torch.no_grad(): output = model(image) loss = criterion(output, target) test_loss += loss.item() count += 1 pred = output.data.cpu().numpy() target = target.cpu().numpy() pred = np.argmax(pred, axis=1) # Add batch sample into evaluator evaluator.add_batch(target, pred) # Fast test during the training test_loss = test_loss / count Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() print('Validation [Epoch: %d, numImages: %5d, Loss: %.3f]' % (epoch, i * batch_size + image.data.shape[0], test_loss)) print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format( Acc, Acc_class, mIoU, FWIoU)) writer.add_scalar('scalar/loss_seg_val', test_loss, epoch) writer.add_scalar('scalar/Acc_val', Acc, epoch) writer.add_scalar('scalar/Acc_class_val', Acc_class, epoch) writer.add_scalar('scalar/mIou_val', mIoU, epoch) writer.add_scalar('scalar/FWIou_val', FWIoU, epoch) path = '/home/user/' if mIoU > best_pred: best_pred = mIoU torch.save(model.state_dict(), path + 'model_params.pkl') print('save the segmentation ' + str(epoch) + 'model, replace the previous parameters')
#%% import numpy as np pred = logit.data.cpu().numpy() target = target.cpu().numpy() pred = np.argmax(pred, axis=1) print(target.shape) #(2, 513, 513) print(pred.shape) #(2, 513, 513) gt_image, pre_image = target, target evaluator = Evaluator(18) evaluator.add_batch(gt_image, pre_image.astype('int')) Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(Acc, Acc_class, mIoU, FWIoU))
class Trainer(object): def __init__(self, args): self.epochs = args.epochs self.save_interval = args.save_interval self.train_data = AerialDataset(args, mode='train') self.train_loader = DataLoader(self.train_data, batch_size=args.train_batch_size, shuffle=True, num_workers=1) self.model = models.deeplabv3_resnet50(num_classes=args.num_of_class) #self.model = models.fcn_resnet50(num_classes=args.num_of_class) self.loss = args.loss if self.loss == 'CE': self.criterion = nn.CrossEntropyLoss() else: #self.loss == 'LS' self.criterion = LovaszSoftmax() self.optimizer = torch.optim.AdamW(self.model.parameters()) self.eval_interval = args.eval_interval self.eval_data = AerialDataset(args, mode='eval') self.eval_loader = DataLoader(self.eval_data, batch_size=args.eval_batch_size, shuffle=False, num_workers=1) self.evaluator = Evaluator(args.num_of_class) self.cuda = args.cuda if self.cuda is True: self.model = self.model.cuda() self.resume = args.resume if self.resume != None: if self.cuda: checkpoint = torch.load(args.resume) else: checkpoint = torch.load(args.resume, map_location='cpu') self.model.load_state_dict(checkpoint['model']) self.optimizer.load_state_dict(checkpoint['opt']) self.start_epoch = checkpoint['epoch'] + 1 #start from next epoch else: self.start_epoch = 1 #Note: self.start_epoch and self.epochs are only used in run() to schedule training & validation def run(self): self.eval(0) end_epoch = self.start_epoch + self.epochs for epoch in range(self.start_epoch, end_epoch): self.train(epoch) if (epoch - self.start_epoch + 1) % self.save_interval == 0: saved_dict = { 'epoch': epoch, 'model': self.model.state_dict(), 'opt': self.optimizer.state_dict() } torch.save(saved_dict, './model_epoch' + str(epoch) + '.pth.tar') if (epoch - self.start_epoch + 1) % self.eval_interval == 0: self.eval(epoch) def train(self, epoch): self.model.train() print(f"----------epoch {epoch}----------") for i, [img, gt] in enumerate(self.train_loader): print("epoch:", epoch, " batch:", i + 1) print("img:", img.shape) print("gt:", gt.shape) self.optimizer.zero_grad() if self.cuda: img, gt = img.cuda(), gt.cuda() pred = self.model(img)['out'] print("pred:", pred.shape) loss = self.criterion(pred, gt.long()) print("loss:", loss) loss.backward() self.optimizer.step() def eval(self, epoch): self.model.eval() self.evaluator.reset() if os.path.exists("epoch" + str(epoch)) is False: os.mkdir("epoch" + str(epoch)) print(f"-----eval epoch {epoch}-----") for i, [ori, img, gt] in enumerate(self.eval_loader): print("batch:", i + 1) print("img:", img.shape) print("gt:", gt.shape) if self.cuda: img = img.cuda() out = self.model(img)['out'] pred = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy() print("pred:", pred.shape) gt = gt.numpy().squeeze() #Both gt and pred are numpy array now self.evaluator.add_batch(gt, pred) #TODO: eval_batch_size > 1 #colorise mask = ret2mask(pred) gt_color = ret2mask(gt) ori = ori.numpy().squeeze().transpose(1, 2, 0) cat = np.concatenate((gt_color, ori, mask), axis=1) cat = Image.fromarray(np.uint8(cat)) cat.save("epoch" + str(epoch) + "/batch" + str(i + 1) + ".png") Acc = self.evaluator.Pixel_Accuracy() Acc_class = self.evaluator.Pixel_Accuracy_Class() mIoU = self.evaluator.Mean_Intersection_over_Union() FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union() print("Acc:", Acc) print("Acc_class:", Acc_class) print("mIoU:", mIoU) print("FWIoU:", FWIoU) def test(self): my_tester = Tester(self) my_tester.test() raise NotImplementedError
def main(): # 配置设备 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 数据路径 train_dir = './data_road/training' valid_dir = './data_road/validing' # 超参数 num_epochs = 20 learning_rate = 0.001 img_size = (640, 192) num_class = 2 SAVE_INTERVAL = 5 evaluator = Evaluator(num_class=num_class) # 构建Dataset实例 train_data = LaneDataset(data_dir=train_dir, img_size=img_size) train_loader = DataLoader(train_data, batch_size=2, shuffle=True) valid_data = LaneDataset(data_dir=valid_dir, img_size=img_size) valid_loader = DataLoader(valid_data, batch_size=2, shuffle=False) # 构建Model实例 # model = FCN8s(n_class=2) # vgg16_path = "models/vgg16_from_caffe.pth" # vgg16 = load_vgg16(model_file=vgg16_path) # model.copy_params_from_vgg16(vgg16) # resnet101-fcn # resnet101_path = './resnet101.pth' # resnet101 = torchvision.models.resnet101(pretrained=False) # state_dict = torch.load(resnet101_path) # resnet101.load_state_dict(state_dict) # model = ResnetFCN(resnet101, num_classes=2, expansion=4) # resnet50-fcn # resnet50_path = './resnet50.pth' # resnet50 = torchvision.models.resnet50(pretrained=False) # state_dict = torch.load(resnet50_path) # resnet50.load_state_dict(state_dict) # model = ResnetFCN(resnet50, num_classes=2, expansion=4) # resnet34-fcn resnet34_path = './resnet34.pth' resnet34 = torchvision.models.resnet34(pretrained=False) state_dict = torch.load(resnet34_path) resnet34.load_state_dict(state_dict) model = ResnetFCN(resnet34, num_classes=2, expansion=1) model = model.to(device) print("模型加载成功!!!") # 定义损失函数和优化器 # criterion = nn.BCELoss() criterion = Focal_Loss() #criterion = nn.BCELoss(weight=torch.tensor([0.3, 0.7])).to(device) #criterion = nn.CrossEntropyLoss(weight=torch.tensor([0.25, 0.75])).to(device) #criterion = nn.CrossEntropyLoss().to(device) optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4) train_Acc_list = [] train_Acc_class_list = [] train_mIoU_list = [] train_FWIoU_list = [] train_loss_list = [] valid_Acc_list = [] valid_Acc_class_list = [] valid_mIoU_list = [] valid_FWIoU_list = [] valid_loss_list = [] # 训练 for epoch in range(num_epochs): train_loss = 0.0 evaluator.reset() for i, (image, label) in enumerate(train_loader): image = image.to(device) label = label.to(device) #print(label[:, 1, :, :]) optimizer.zero_grad() # 前向传播 output = model(image) loss = criterion(output, label) #output = torch.sigmoid(output) output = torch.softmax(output, dim=1) #loss = criterion(output.transpose(1,3), label.transpose(1,3)) output = torch.argmax(output, dim=1).cpu().numpy() label = torch.argmax(label, dim=1).cpu().numpy() #print(output) # 后向传播及优化 loss.backward() train_loss += loss.item() optimizer.step() evaluator.add_batch(label, output) # 添加output和label,用于后续评估 # 计算像素准确率、平均像素准确率、平均交并比、频权交并比 Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() train_Acc_list.append(Acc) train_Acc_class_list.append(Acc_class) train_mIoU_list.append(mIoU) train_FWIoU_list.append(FWIoU) train_loss_list.append(train_loss / len(train_loader)) evaluator.reset() # 保存模型 if (epoch + 1) == num_epochs: #epoch % SAVE_INTERVAL == 0 or torch.save(model.state_dict(), './models/fcn8s_BFL.pth') print( "Epoch_{}: train_loss: {:.6f} Acc: {:.4f} Acc_class: {:.4f} mIoU: {:.4f} FWIoU: {:.4f}" .format(epoch + 1, train_loss / len(train_loader), Acc, Acc_class, mIoU, FWIoU)) # 验证阶段 model.eval() valid_loss = 0.0 with torch.no_grad(): for i, (image, label) in enumerate(valid_loader): image = image.to(device) label = label.to(device) output = model(image) loss = criterion(output, label) #output = torch.sigmoid(output) output = torch.softmax(output, dim=1) # loss = criterion(output.transpose(1, 3), label.transpose(1, 3)) output = torch.argmax(output, dim=1).cpu().numpy() label = torch.argmax(label, dim=1).cpu().numpy() #print(output) #print(label) valid_loss += loss.item() evaluator.add_batch(label, output) # 添加output和label,用于后续评估 # 计算像素准确率、平均像素准确率、平均交并比、频权交并比 Acc = evaluator.Pixel_Accuracy() Acc_class = evaluator.Pixel_Accuracy_Class() mIoU = evaluator.Mean_Intersection_over_Union() FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union() print( "Epoch_{}: valid_loss: {:.6f} Acc: {:.4f} Acc_class: {:.4f} mIoU: {:.4f} FWIoU: {:.4f}" .format(epoch + 1, valid_loss / len(valid_loader), Acc, Acc_class, mIoU, FWIoU)) valid_Acc_list.append(Acc) valid_Acc_class_list.append(Acc_class) valid_mIoU_list.append(mIoU) valid_FWIoU_list.append(FWIoU) valid_loss_list.append(valid_loss / len(valid_loader)) # 绘制曲线 draw_figures(train_loss_list, valid_loss_list, title='loss', y_label='loss') draw_figures(train_Acc_list, valid_Acc_list, title='Acc', y_label='Acc') draw_figures(train_Acc_class_list, valid_Acc_class_list, title='Acc_class', y_label='Acc_class') draw_figures(train_mIoU_list, valid_mIoU_list, title='mIoU', y_label='mIoU') draw_figures(train_FWIoU_list, valid_FWIoU_list, title='FWIoU', y_label='FWIoU') print("完成曲线绘制!!!")
def train(data_set_type, num_classes, batch_size, epochs, use_gpu, learning_rate, w_decay): model = get_fcn_model(num_classes, use_gpu) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(params=model.parameters(), lr=learning_rate, weight_decay=w_decay) scheduler = lr_scheduler.StepLR( optimizer, step_size=step_size, gamma=gamma) # decay LR by a factor of 0.5 every 5 epochs data_set, data_loader = get_dataset_dataloader(data_set_type, batch_size) since = time.time() best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 epoch_loss = np.zeros((2, epochs)) epoch_acc = np.zeros((2, epochs)) epoch_iou = np.zeros((2, epochs, num_classes)) epoch_mean_iou = np.zeros((2, epochs)) evaluator = Evaluator(num_classes) for epoch in range(epochs): logger.info('Epoch {}/{}'.format(epoch + 1, epochs)) logger.info('-' * 28) for phase_ind, phase in enumerate(['train', 'val']): if phase == 'train': model.train() logger.info(phase) else: model.eval() logger.info(phase) evaluator.reset() running_loss = 0.0 running_acc = 0.0 num_of_batches = math.ceil(len(data_set[phase]) / batch_size) running_iou = np.zeros((num_of_batches, num_classes)) for batch_ind, batch in enumerate(data_loader[phase]): imgs, targets = batch imgs = Variable(imgs).float() imgs = imgs.to(device) targets = Variable(targets).type(torch.LongTensor) targets = targets.to(device) # zero the learnable parameters gradients optimizer.zero_grad() with torch.set_grad_enabled(phase == 'train'): outputs = model(imgs) loss = criterion(outputs, targets) if phase == 'train': loss.backward() optimizer.step() # computes loss and acc for current iteration preds = torch.argmax(outputs, dim=1) ious = iou(preds, targets, num_classes) running_loss += loss * imgs.size(0) running_acc += pixelwise_acc(preds, targets) * imgs.size(0) running_iou[batch_ind, :] = ious logger.debug('Batch {} running loss: {}'.format( batch_ind, running_loss)) # test the iou and pixelwise accuracy using evaluator preds = preds.cpu().numpy() targets = targets.cpu().numpy() evaluator.add_batch(targets, preds) epoch_loss[phase_ind, epoch] = running_loss / len(data_set[phase]) epoch_acc[phase_ind, epoch] = running_acc / len(data_set[phase]) epoch_iou[phase_ind, epoch] = np.nanmean(running_iou, axis=0) epoch_mean_iou[phase_ind, epoch] = np.nanmean(epoch_iou[phase_ind, epoch]) logger.info('{} loss: {:.4f}, acc: {:.4f}, mean iou: {:.6f}'.format(phase,\ epoch_loss[phase_ind, epoch], epoch_acc[phase_ind, epoch],\ epoch_mean_iou[phase_ind, epoch])) eva_pixel_acc = evaluator.Pixel_Accuracy() eva_pixel_acc_class = evaluator.Pixel_Accuracy_Class() eva_mIOU = evaluator.Mean_Intersection_over_Union() logger.info('{} - Evaluator - acc: {:.4f}, acc class: {:.4f}, mean iou: {:.6f}'.format(phase,\ eva_pixel_acc, eva_pixel_acc_class, eva_mIOU)) if phase == 'val' and epoch_acc[phase_ind, epoch] > best_acc: best_acc = epoch_acc[phase_ind, epoch] best_model_wts = copy.deepcopy(model.state_dict()) print() time_elapsed = time.time() - since logger.info('Training completed in {}m {}s'.format(int(time_elapsed / 60),\ int(time_elapsed) % 60)) # load best model weights model.load_state_dict(best_model_wts) # save numpy results np.save(os.path.join(score_dir, 'epoch_accuracy'), epoch_acc) np.save(os.path.join(score_dir, 'epoch_mean_iou'), epoch_mean_iou) np.save(os.path.join(score_dir, 'epoch_iou'), epoch_iou) return model