From 60f74e3e41254c0c8f5864ec5b601542014ae2b9 Mon Sep 17 00:00:00 2001 From: xufeifei Date: Wed, 27 Mar 2019 17:06:55 +0800 Subject: [PATCH] add ignored_Id --- train/loadData.py | 16 ++++++++++++---- train/main.py | 21 ++++++++++++++------- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/train/loadData.py b/train/loadData.py index c1659c2..ed54c76 100644 --- a/train/loadData.py +++ b/train/loadData.py @@ -8,7 +8,7 @@ class LoadData: ''' Class to laod the data ''' - def __init__(self, data_dir, classes, cached_data_file, normVal=1.10): + def __init__(self, data_dir, classes, cached_data_file, ignored_Id, normVal=1.10): ''' :param data_dir: directory where the dataset is kept :param classes: number of classes in the dataset @@ -26,6 +26,7 @@ def __init__(self, data_dir, classes, cached_data_file, normVal=1.10): self.trainAnnotList = list() self.valAnnotList = list() self.cached_data_file = cached_data_file + self.ignored_Id = ignored_Id def compute_class_weights(self, histogram): ''' @@ -85,9 +86,16 @@ def readFile(self, fileName, trainStg=False): self.valAnnotList.append(label_file) if max_val > (self.classes - 1) or min_val < 0: - print('Labels can take value between 0 and number of classes.') - print('Some problem with labels. Please check.') - print('Label Image ID: ' + label_file) + if max_val == self.ignored_Id: + print('Label id: %d has been ignored' %self.ignored_Id) + print('Label Image ID: ' + label_file) + + else: + print('Labels can take value between 0 and number of classes.') + print('Some problem with labels.' + 'You might want to set ignored_Id = ' + 'Please check argument ignored_Id in main.py') + print('Label Image ID: ' + label_file) no_files += 1 if trainStg == True: diff --git a/train/main.py b/train/main.py index 50d3858..98604e1 100644 --- a/train/main.py +++ b/train/main.py @@ -48,14 +48,16 @@ def val(args, val_loader, model, criterion): # compute the loss loss = criterion(output, target_var) - epoch_loss.append(loss.data[0]) + #epoch_loss.append(loss.data[0]) + + epoch_loss.append(loss.item()) time_taken = time.time() - start_time # compute the confusion matrix iouEvalVal.addBatch(output.max(1)[1].data, target_var.data) - print('[%d/%d] loss: %.3f time: %.2f' % (i, total_batches, loss.data[0], time_taken)) + print('[%d/%d] loss: %.3f time: %.2f' % (i, total_batches, loss.item(), time_taken)) average_epoch_loss_val = sum(epoch_loss) / len(epoch_loss) @@ -102,13 +104,13 @@ def train(args, train_loader, model, criterion, optimizer, epoch): loss.backward() optimizer.step() - epoch_loss.append(loss.data[0]) + epoch_loss.append(loss.item()) time_taken = time.time() - start_time #compute the confusion matrix iouEvalTrain.addBatch(output.max(1)[1].data, target_var.data) - print('[%d/%d] loss: %.3f time:%.2f' % (i, total_batches, loss.data[0], time_taken)) + print('[%d/%d] loss: %.3f time:%.2f' % (i, total_batches, loss.item(), time_taken)) average_epoch_loss_train = sum(epoch_loss) / len(epoch_loss) @@ -149,7 +151,7 @@ def trainValidateSegmentation(args): ''' # check if processed data file exists or not if not os.path.isfile(args.cached_data_file): - dataLoad = ld.LoadData(args.data_dir, args.classes, args.cached_data_file) + dataLoad = ld.LoadData(args.data_dir, args.classes, args.cached_data_file, args.ignored_Id) data = dataLoad.processData() if data is None: print('Error while pickling data. Please check.') @@ -159,6 +161,7 @@ def trainValidateSegmentation(args): q = args.q p = args.p + ignored_Id = args.ignored_Id # load the model if not args.decoder: model = net.ESPNet_Encoder(args.classes, p=p, q=q) @@ -192,7 +195,8 @@ def trainValidateSegmentation(args): if args.onGPU: weight = weight.cuda() - criteria = CrossEntropyLoss2d(weight) #weight + #criteria = CrossEntropyLoss2d(weight) #weight + criteria = torch.nn.CrossEntropyLoss(weight=weight,ignore_index=ignored_Id) #ignore index -100 for default if args.onGPU: criteria = criteria.cuda() @@ -386,7 +390,7 @@ def trainValidateSegmentation(args): parser.add_argument('--scaleIn', type=int, default=8, help='For ESPNet-C, scaleIn=8. For ESPNet, scaleIn=1') parser.add_argument('--max_epochs', type=int, default=300, help='Max. number of epochs') parser.add_argument('--num_workers', type=int, default=4, help='No. of parallel threads') - parser.add_argument('--batch_size', type=int, default=12, help='Batch size. 12 for ESPNet-C and 6 for ESPNet. ' + parser.add_argument('--batch_size', type=int, default=8, help='Batch size. 12 for ESPNet-C and 6 for ESPNet. ' 'Change as per the GPU memory') parser.add_argument('--step_loss', type=int, default=100, help='Decrease learning rate after how many epochs.') parser.add_argument('--lr', type=float, default=5e-4, help='Initial learning rate') @@ -402,6 +406,9 @@ def trainValidateSegmentation(args): 'Only used when training ESPNet') parser.add_argument('--p', default=2, type=int, help='depth multiplier') parser.add_argument('--q', default=8, type=int, help='depth multiplier') + parser.add_argument('--ignored_Id', default=255, type=int, help='ignoredTrainId for crossEntryLoss.' + 'default 255 to ignore cityscapes background TrainId.' + 'the background Id could be seen in cityscapesScripts/cityscapesscripts/helpers/label.py') trainValidateSegmentation(parser.parse_args())