forked from melissande/dhi-segmentation-buildings
-
Notifications
You must be signed in to change notification settings - Fork 0
/
train_model.py
417 lines (310 loc) · 16.2 KB
/
train_model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
import sys
import numpy as np
import os
import cv2
import matplotlib.pyplot as plt
import logging
import torch
import torch.nn as nn
from torch import optim
import torch.backends.cudnn as cudnn
from torchvision import transforms
from torch.utils.data import DataLoader
import time
from random import randint
from IOU_computations import *
from Data_Handle.dataset_generator import Dataset_sat
from predict_and_evaluate import *
from Data_Handle.data_augmentation import *
###### PATH TO STORE MODEL ############
GLOBAL_PATH='MODEL_TEST_GHANA/'
if not os.path.exists(GLOBAL_PATH):
os.makedirs(GLOBAL_PATH)
######################################
###################################
INPUT_CHANNELS=9 #9 channels for panchromatic + 8 pansharpened. If not set to 9, plotting of patches will mess up.
# so only works for INPUT_CHANNELS=9 anyway.
NB_CLASSES=2 #Building and Background. Only works for NB_CLASSES=2 anyway, otherwise this network doesn't work.
SIZE_PATCH=128# patches of size 128x128. Needs to be equal to the size of the patches of the dataset.
##############
MODEL_PATH_SAVE=GLOBAL_PATH+'RESUNET_test_'
MODEL_PATH_RESTORE='' #Path of Model to restore ex: 'TRAINED_MODELS/RUBV3D2_final_model_ghana.pth'
TEST_SAVE=GLOBAL_PATH+'TEST_SAVE/' #to store some patches initial and final epoch of validation set + models + performance curves
if not os.path.exists(TEST_SAVE):
os.makedirs(TEST_SAVE)
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
##############
DROPOUT=0.35
DEFAULT_BATCH_SIZE =8 #has to be set to 8 for Ghana dataset and can be set up to 32 for Spacenet dataset
DEFAULT_EPOCHS =2
DEFAULT_VALID=32 # Batch size for validation set.
#Knowing that around 1200 elements in ghana validation and 15000 in spacenet validation
DISPLAY_STEP=100 #how often (in terms of iterations) is displayed measures during an epoch
IOU_STEP=15 # how often is computed IOU measures over validations et
###############
DEFAULT_LAYERS=3 #number of layers of the UNET (not considering bottom layer) = number of downsmapling stages
DEFAULT_FEATURES_ROOT=32 # number of filters in the first layer of the Unet
DEFAULT_BN=True # Batch normalization layers included
#####
DEFAULT_FILTER_WIDTH=3 #convolution kernel size. ex, here: 3x3
DEFAULT_LR=1e-3#1e-3for spacenet and ghana
DEFAULT_N_RESBLOCKS=1 #can add residual blocks inside each stage. Make the network heavier. Not advised.
###Tune Learning rate
REDUCE_LR_STEPS = [1,5, 50, 100,200] #reduce everytime one of these epochs is reached
################
DISTANCE_NET='v2' #can be set to none if no distance module wants to be used
BINS=10
THRESHOLD=20
if DISTANCE_NET is None:
DISTANCE_NET_UNET=False # has to be set to False if no distance module is used, otherwise error.
else:
DISTANCE_NET_UNET=True
##### Data augmentation set for training ###
DATA_AUG=None
# DATA_AUG=transforms.Compose([Transform(),ToTensor()])
####### TMP folder for IOU ###
## not to worry about, compulsory for vectorizing masks ""
####### Data: where the dataset is stored ###
# root_folder ='../SPACENET_DATA/SPACENET_DATA_PROCESSED/DATASET/128_x_128_8_bands_pansh/'
root_folder = '../2_DATA_GHANA/DATASET/128_x_128_8_pansh/'
#type of loss used
LOSS_FN='cross-entropy'# or 'jaccard_approx'
class Trainer(object):
"""
Trains a unet instance
:param net: the unet instance to train
:param batch_size: size of training batch
:param lr: learning rate
:nb_classes: always set to 2 ->background and building
:type of loss: 'cross-entropy' or 'jaccard_approx-approx'
"""
def __init__(self, net, batch_size=32, lr=0.001,nb_classes=2,loss_fn=LOSS_FN):
self.net = net
self.batch_size = batch_size
self.lr = lr
self.nb_classes=nb_classes
self.loss_fn=loss_fn
def _initialize(self, prediction_path,store_learning,iou_step,dist_net,threshold,bins):
self.optimizer = optim.Adam(self.net.parameters(),lr=self.lr)
self.prediction_path = prediction_path
self.store_learning=store_learning
self.IOU_STEP=iou_step
self.threshold=threshold
self.bins=bins
self.dist_net=dist_net
def train(self, data_provider_path,store_learning, save_path='', restore_path='', epochs=3, dropout=0.2, display_step=100, validation_batch_size=30, prediction_path = '',dist_net=None,threshold=20,bins=15,iou_step=1,reduce_lr_steps=[1,10,100,200],data_aug=None):
"""
Lauches the training process
:param data_provider_path: where the DATASET folder is
:param store_learning: to store the metrics during the training as .txt file
:param save_path: path where to store checkpoints
:param restore_path: path where is the model to restore is stored
:param epochs: number of epochs
:param dropout: dropout probability
:param validation_batch_size: batch size of the validation set
:param prediction_path: where to store output of training (patches, losses .txt file, models)
:param dist_net: distance module or not
:param threshold: threshold of distance module
:param bins: number of bins for distance module
:iou_step: how often is computed Iou measures over the validation set
:reduce_lr_steps: epoch at which the learning rate is halved
:data_aug: 'yes' or 'no' if the training set is augmented
"""
##SET UP PATHS FOR TRAINING ##
#check they exist?
PATH_TRAINING=data_provider_path+'TRAINING/'
if not os.path.exists(PATH_TRAINING):
print('Training dataset path not valid. Should be path_to_dataset/TRAINING/ and this folder should contain INTPUT/ and OUTPUT/')
raise
PATH_VALIDATION=data_provider_path+'VALIDATION/'
if not os.path.exists(PATH_VALIDATION):
print('Validation dataset path not valid. Should be path_to_dataset/VALIDATION/ and this folder should contain INTPUT/ and OUTPUT/')
raise
PATH_TEST=data_provider_path+'TEST/'
if not os.path.exists(PATH_TEST):
print('Test dataset path not valid. Should be path_to_dataset/TEST/ and this folder should contain INTPUT/ and OUTPUT/')
raise
TMP_IOU=prediction_path+'TMP_IOU/'
if not os.path.exists(TMP_IOU):
os.makedirs(TMP_IOU)
loss_train=[]
if epochs == 0:
print('Epoch set 0, model won\'t be trained')
raise
if save_path=='':
print('Specify a path where to store the Model')
raise
if prediction_path=='':
print('Specify where to stored visualization of training')
raise
if restore_path=='':
store_learning.initialize('w')
store_learning
print('Model trained from scratch')
else:
store_learning.initialize('a')
self.net.load_state_dict(torch.load(restore_path))
print('Model loaded from {}'.format(restore_path))
self._initialize(prediction_path,store_learning,iou_step,dist_net,threshold,bins)
###Validation loader
val_generator=Dataset_sat.from_root_folder(PATH_VALIDATION,self.nb_classes)
val_loader = DataLoader(val_generator, batch_size=validation_batch_size,shuffle=False, num_workers=1)
RBD=randint(0,int(val_loader.__len__())-1)
self.info_validation(val_loader,-1,RBD,"_init",TMP_IOU)
###Training loader
train_generator=Dataset_sat.from_root_folder(PATH_TRAINING,self.nb_classes,transform=data_aug)#max_data_size=4958
logging.info("Start optimization")
counter=0
for epoch in range(epochs):
##tune learning reate
if epoch in reduce_lr_steps:
self.lr = self.lr * 0.5
self.optimizer = torch.optim.Adam(self.net.parameters(), lr=self.lr)
total_loss = 0
error_tot=0
train_loader = DataLoader(train_generator, batch_size=self.batch_size,shuffle=True, num_workers=1)
for i_batch,sample_batch in enumerate(train_loader):
self.optimizer.zero_grad()
predict_net=Train_or_Predict(sample_batch,self.dist_net,self.loss_fn,self.threshold,self.bins,self.net)
loss,_,probs_seg=predict_net.forward_pass()
loss,self.optimizer,self.net=predict_net.backward_prog(loss,self.optimizer)
total_loss+=loss.data[0]
loss_train.append(loss.data[0])
counter+=1
if i_batch % display_step == 0:
self.output_training_stats(i_batch,loss,predict_net.batch_y,probs_seg)
avg_loss_train_value=total_loss/train_loader.__len__()
(self.store_learning).avg_loss_train.append(avg_loss_train_value)
(self.store_learning).write_file((self.store_learning).file_train,avg_loss_train_value)
logging.info(" Training {:}, Minibatch Loss= {:.4f}".format("epoch_%s"%epoch,avg_loss_train_value))
self.info_validation(val_loader,epoch,RBD,"epoch_%s"%epoch,TMP_IOU)
torch.save(self.net.state_dict(),save_path + 'CP{}.pth'.format(epoch))
print('Checkpoint {} saved !'.format(epoch))
self.info_validation(val_loader,-2,RBD,'_last_',TMP_IOU)
# time.sleep(4)
# plt.close(fig)
return save_path + 'CP{}.pth'.format(epoch)
def output_training_stats(self, step, loss,batch_y,probs_seg):
# Calculate batch loss and accuracy
loss_v=loss.data[0]
groundtruth_seg_v=np.asarray(batch_y)
prediction_seg_v=probs_seg.data.cpu().numpy()
logging.info("Iter {:}, Minibatch Loss= {:.4f}, Minibatch error= {:.4f}%".format(step,loss_v,error_rate(prediction_seg_v, groundtruth_seg_v)))
def info_validation(self,val_loader,epoch,RBD,name,TMP_IOU):
loss_v=0
error_rate_v=0
iou_acc_v=0
f1_v=0
if name=="_init":
display_patches=True
save_patches=True
save_IOU_metrics=False
elif name=='_last_':
display_patches=True
save_patches=True
save_IOU_metrics=False
else:
display_patches=True
save_patches=False
save_IOU_metrics=True
for i_batch,sample in enumerate(val_loader):
predict_net=Train_or_Predict(sample,self.dist_net,self.loss_fn,self.threshold,self.bins,self.net)
loss,probs_dist,probs_seg=predict_net.forward_pass()
prediction_seg_v=probs_seg.data.cpu().numpy()
groundtruth_seg_v=np.asarray(predict_net.batch_y)
prediction_dist_v=probs_dist.data.cpu().numpy()
groundtruth_dist=np.asarray(predict_net.batch_y_dist)
loss_v+=loss.data[0]
error_rate_v+=error_rate(prediction_seg_v,groundtruth_seg_v)
if (save_IOU_metrics and (epoch+1)%self.IOU_STEP==0):
iou_acc,f1,_=predict_score_batch(TMP_IOU,np.argmax(groundtruth_seg_v,3),np.argmax(prediction_seg_v,3))
iou_acc_v+=iou_acc
f1_v+=f1
loss_v/=val_loader.__len__()
error_rate_v/=val_loader.__len__()
logging.info("Verification loss= {:.4f},error= {:.4f}%".format(loss_v,error_rate_v))
if (name!="_init" and name!='_last_'):
(self.store_learning).write_file((self.store_learning).file_verif,loss_v)
(self.store_learning).write_file((self.store_learning).error_rate_file_verif,error_rate_v)
if (save_IOU_metrics and (epoch+1)%self.IOU_STEP==0):
iou_acc_v/=val_loader.__len__()
f1_v/=val_loader.__len__()
logging.info("Verification IOU Precision = {:.4f}%, F1 IOU= {:.4f}%".format(iou_acc_v,f1_v))
(self.store_learning).write_file((self.store_learning).IOU_acc_file_verif,iou_acc_v)
(self.store_learning).write_file((self.store_learning).f1_IOU_file_verif,f1_v)
if __name__ == '__main__':
# python train_model.py ../2_DATA_GHANA/DATASET/128_x_128_8_pansh/ MODEL_TEST_GHANA/ RESUNET_test_ '' --epochs=6 --iou_step=2
root_folder=sys.argv[1]
##########
GLOBAL_PATH=sys.argv[2]
if not os.path.exists(GLOBAL_PATH):
os.makedirs(GLOBAL_PATH)
TEST_SAVE=GLOBAL_PATH+'TEST_SAVE/'
if not os.path.exists(TEST_SAVE):
os.makedirs(TEST_SAVE)
##########
MODEL_PATH_SAVE=GLOBAL_PATH+sys.argv[3]
MODEL_PATH_RESTORE=sys.argv[4]
for i in range(5, len(sys.argv)):
arg = sys.argv[i]
if arg.startswith('--input_channels'):
INPUT_CHANNELS=int(arg[len('--input_channels='):])
elif arg.startswith('--nb_classes'):
NB_CLASSES=int(arg[len('--nb_classes='):])
# elif arg.startswith('--unet_version'):
# UNET_V=int(arg[len('--unet_version='):])
elif arg.startswith('--nb_layers'):
DEFAULT_LAYERS=int(arg[len('--nb_layers='):])
elif arg.startswith('--filter_width'):
DEFAULT_FILTER_WIDTH=int(arg[len('--filter_width='):])
elif arg.startswith('--nb_features_root'):
DEFAULT_FEATURES_ROOT=int(arg[len('--nb_features_root='):])
elif arg.startswith('--learning_rate'):
DEFAULT_LR=float(arg[len('--learning_rate='):])
elif arg.startswith('--batch_size'):
DEFAULT_BATCH_SIZE = int(arg[len('--batch_size='):])
elif arg.startswith('--epochs'):
DEFAULT_EPOCHS = int(arg[len('--epochs='):])
elif arg.startswith('--dropout'):
DROPOUT = float(arg[len('--dropout='):])
elif arg.startswith('--display_step'):
DISPLAY_STEP = int(arg[len('--display_step='):])
elif arg.startswith('--validation_size_batch'):
DEFAULT_VALID = int(arg[len('--validation_size_batch='):])
elif arg.startswith('--distance_net'):
DISTANCE_NET = arg[len('--distance_net='):]
if DISTANCE_NET=='v2':
BINS=10
THRESHOLD=20
DISTANCE_NET_UNET=True
elif DISTANCE_NET=='None':
DISTANCE_NET=None
DISTANCE_NET_UNET=False
else:
raise ValueError('Unknown argument %s' % str(arg))
elif arg.startswith('--batch_norm'):
DEFAULT_BN = eval(arg[len('--batch_norm='):])
elif arg.startswith('--iou_step'):
IOU_STEP = int(arg[len('--iou_step='):])
elif arg.startswith('--lr_reduce_steps'):
REDUCE_LR_STEPS = np.asarray(arg[len('--lr_reduce_steps='):].split(',')).astype(int)
elif arg.startswith('--data_aug'):
if (arg[len('--data_aug='):].lower()=='yes'):
DATA_AUG=transforms.Compose([Transform(),ToTensor()])
elif (arg[len('--data_aug='):].lower()=='no'):
DATA_AUG=None
else:
raise ValueError('Unknown argument %s' % str(arg))
elif arg.startswith('--loss_func'):
LOSS_FN=arg[len('--loss_func='):]
else:
raise ValueError('Unknown argument %s' % str(arg))
from RUBV3D2 import UNet
model=UNet(INPUT_CHANNELS,NB_CLASSES,depth =DEFAULT_LAYERS,n_features_zero =DEFAULT_FEATURES_ROOT,width_kernel=DEFAULT_FILTER_WIDTH,dropout=DROPOUT,distance_net=DISTANCE_NET_UNET,bins=BINS,batch_norm=DEFAULT_BN)
model.cuda()
cudnn.benchmark = True
print('### Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0))
trainer=Trainer(model,DEFAULT_BATCH_SIZE,DEFAULT_LR,NB_CLASSES,LOSS_FN)
store_learning=Store_learning(GLOBAL_PATH)
save_path=trainer.train( root_folder,store_learning, MODEL_PATH_SAVE, MODEL_PATH_RESTORE,DEFAULT_EPOCHS,DROPOUT, DISPLAY_STEP, DEFAULT_VALID, TEST_SAVE,DISTANCE_NET,THRESHOLD,BINS,IOU_STEP,REDUCE_LR_STEPS,DATA_AUG)
print('Last model saved is %s: '%save_path)