def validate(model, valloader, n_class):
import pdb;pdb.set_trace()
losses = AverageMeter()
model.eval()
gts, preds = [], []
for i, (images, labels) in enumerate(valloader):
images = Variable(images.cuda())
labels = Variable(labels.cuda())
outputs = model(images)
if(isinstance(outputs, tuple)):
outputs = outputs[0]
loss = cross_entropy2d(outputs, labels)
losses.update(loss.data[0], images.size(0))
gt = labels.data.cpu().numpy()
pred = outputs.data.max(1)[1].cpu().numpy()
#pred = outputs.data[:,1:,:,:].max(1)[1].cpu().numpy() + 1
for gt_, pred_ in zip(gt, pred):
gts.append(gt_)
preds.append(pred_)
score = scores(gts, preds, n_class=n_class)
return losses.avg, score
def train(args):
# Setup train DataLoader
trainloader = CCFLoader(args.traindir, split=args.split,
is_transform=True, img_size=(args.img_rows, args.img_cols))
n_classes = trainloader.n_classes
TrainDataLoader = data.DataLoader(
trainloader, batch_size=args.batch_size, num_workers=4, shuffle=True)
# Setup validate DataLoader
valloader = CCFLoader(args.traindir, split='val', is_transform=True, img_size=(
args.img_rows, args.img_cols))
VALDataLoader = data.DataLoader(
valloader, batch_size=4, num_workers=4, shuffle=False)
# Setup visdom for visualization
vis = visdom.Visdom()
assert vis.check_connection()
loss_window = vis.line(X=np.zeros((1,)),
Y=np.zeros((1)),
opts=dict(xlabel='minibatches',
ylabel='Loss',
title=args.arch+' Training Loss',
legend=['Loss']))
valacc_window = vis.line(X=np.zeros((1,)),
Y=np.zeros((1)),
opts=dict(xlabel='minibatches',
ylabel='ACC',
title='Val ACC',
legend=['ACC']))
# Setup model
if(args.snapshot == None):
model = get_model(args.arch, n_classes)
model = DataParallel(model.cuda(args.gpu[0]), device_ids=args.gpu)
start_epoch = 0
else:
model = get_model(args.arch, n_classes)
state_dict = torch.load(args.snapshot).state_dict()
from collections import OrderedDict
new_state_dict = OrderedDict()
for key, value in list(state_dict.items()):
original_key = key[7:] # remove 'moudle.'
new_state_dict[original_key] = value
model.load_state_dict(new_state_dict)
model = DataParallel(model.cuda(), device_ids=[i for i in range(len(args.gpu))])
start_epoch = int(os.path.basename(args.snapshot).split('.')[0])
optimizer = torch.optim.SGD(model.parameters(), lr=args.l_rate, momentum=0.99, weight_decay=5e-4)
print(model)
# Start training
for epoch in range(args.n_epoch):
adjust_learning_rate(optimizer, args.l_rate, epoch, args.step)
if(epoch < start_epoch):
continue
print("Epoch [%d/%d] learning rate: %f" % (epoch+1, args.n_epoch, optimizer.param_groups[0]['lr']))
for i, (images, labels) in enumerate(TrainDataLoader):
if torch.cuda.is_available():
images = Variable(images.cuda(args.gpu[0]))
labels = Variable(labels.cuda(args.gpu[0]))
else:
images = Variable(images)
labels = Variable(labels)
iter = len(TrainDataLoader)*epoch + i
#poly_lr_scheduler(optimizer, args.l_rate, iter)
model.train()
optimizer.zero_grad()
outputs = model(images)
if(isinstance(outputs, tuple)):
loss = cross_entropy2d(outputs[0], labels, weights_per_class) + args.clsloss_weight * bin_clsloss(outputs[1], labels)
else:
#loss = cross_entropy2d(outputs, labels)
loss = cross_entropy2d(outputs, labels, weights_per_class)
#loss = focal_loss2d(outputs, labels)
loss.backward()
optimizer.step()
vis.line(
X=torch.ones((1, 1)).cpu()*iter,
Y=torch.Tensor([loss.data]).unsqueeze(0).cpu(),
win=loss_window,
update='append')
print("Epoch [%d/%d] loss: %f" % (epoch+1, args.n_epoch, loss))
# validation
loss, score = validate(model, VALDataLoader, n_classes)
for i in range(n_classes):
print(i, score['Class Acc'][i])
vis.line(
X=torch.ones((1, 1)).cpu()*(epoch+1),
Y=torch.ones((1, 1)).cpu()*score['Overall Acc'],
win=valacc_window,
update='append')
if(not os.path.exists("snapshot/{}".format(args.arch))):
os.mkdir("snapshot/{}".format(args.arch))
torch.save(model, "snapshot/{}/{}.pkl".format(args.arch, epoch+1))
def main():
print('RUNDIR: {}'.format(args.logdir))
sys.stdout.flush()
logger = function.get_logger(args.logdir)
logger.info('test') # write in log file
running_metrics_val = function.runningScoreSeg(args.n_class)
val_loss_meter = function.averageMeter()
time_meter = function.averageMeter()
print(len(valloader))
start_ts = time.time() # return current time stamp
model.eval()
with torch.no_grad():
for i_val, (leftimgval, rightimgval, labelval, disp_true,
L_name) in tqdm(enumerate(valloader)):
imgL = leftimgval.numpy()
imgR = rightimgval.numpy()
if args.cuda:
imgL = torch.FloatTensor(imgL).cuda()
imgR = torch.FloatTensor(imgR).cuda()
imgL, imgR = Variable(imgL), Variable(imgR)
output = model(imgL, imgR) # 1 1024 2048 1 19 1024 2048
pred_seg = output0.data.cpu().numpy()
# FCN OF SEGMETATION
N, _, h, w = pred_seg.shape # 4,12,192,704,numpy
pred_segmap = pred_seg.transpose(0, 2, 3, 1).reshape(
-1, args.n_class).argmax(axis=1).reshape(N, h, w)
img = drawseg.direct_render(pred_segmap, args.n_class)
skimage.io.imsave(args.saveseg + (L_name[0].split('/')[-1]),
img[0])
# segmetation mIoU
score = torch.from_numpy(pred_seg).cuda()
lossval = cross_entropy2d(
score, labelval.cuda()) # mean pixelwise loss in a batch
pred = score.data.max(
1)[1].cpu().numpy() # [batch_size, height, width]#229,485
gt = labelval.data.cpu().numpy(
) # [batch_size, height, width]#256,512
running_metrics_val.update(gt=gt, pred=pred)
val_loss_meter.update(lossval.item())
torch.cuda.empty_cache()
logger.info(" val_loss: %.4f" % (val_loss_meter.avg))
print("val_loss: %.4f" % (val_loss_meter.avg))
#"""
# output scores
score, class_iou = running_metrics_val.get_scores()
for k, v in score.items():
print(k, v)
sys.stdout.flush()
logger.info('{}: {}'.format(k, v))
for k, v in class_iou.items():
print(k, v)
logger.info('{}: {}'.format(k, v))
def main():
writer = SummaryWriter(log_dir=args.logdir)
print('RUNDIR: {}'.format(args.logdir))
sys.stdout.flush()
#shutil.copy(args.config, args.logdir) # copy config file to path of logdir
logger = function.get_logger(args.logdir)
logger.info('Let the games begin') # write in log file
running_metrics_val = function.runningScoreSeg(args.n_class)
val_loss_meter = function.averageMeter()
time_meter = function.averageMeter()
best_iou = -100.0
i = 0
flag = True
train_iters=200000
print_iter= 150 #150
val_iters=200 #750
save_iters=10000
while i <= train_iters and flag: # 3000000
print(len(trainloader))
for (leftimg,labels) in trainloader:
start_ts = time.time() # return current time stamp
model.train() # set model to training mode
imgL = Variable(leftimg).cuda()
label = Variable(labels).cuda()#1,256,512
optimizer.zero_grad() # clear earlier gradients
output0 = model(imgL)
loss = cross_entropy2d(output0, label).to(device)
loss.backward() # backpropagation loss
optimizer.step() # optimizer parameter update
time_meter.update(time.time() - start_ts)
if (i + 1) %print_iter == 0: # 150
fmt_str = "Iter [{:d}/{:d}] Loss: {:.4f} Time/Image: {:.4f}"
print_str = fmt_str.format(i + 1,train_iters,loss.item(),time_meter.val / args.batch_size)
print(print_str)
logger.info(print_str)
writer.add_scalar('loss/train_loss', loss.item(), i + 1 )
time_meter.reset()
if (i + 1) % save_iters == 0:
savefilename = args.savemodel + 'iter_' + str(i) + '.tar'
torch.save({
'iter': i,
'state_dict': model.state_dict(),
}, savefilename)
torch.cuda.empty_cache()
if (i + 1) % val_iters == 0 or (i + 1) == train_iters: # 750
model.eval()
with torch.no_grad():
for i_val, (leftimgval, labelval) in tqdm(enumerate(valloader)):
imgL = Variable(leftimgval.cuda())
output0val = model(imgL) # [batch_size, n_classes, height, width]
#segmetation
lossval = cross_entropy2d(output0val, labelval.cuda())# mean pixelwise loss in a batch
pred = output0val.data.max(1)[1].cpu().numpy() # [batch_size, height, width]#229,485
gt = labelval.data.cpu().numpy() # [batch_size, height, width]#256,512
running_metrics_val.update(gt=gt, pred=pred)
val_loss_meter.update(lossval.item())
writer.add_scalar('loss/val_loss', val_loss_meter.avg, i + 1)
logger.info("Iter %d val_loss: %.4f" % (i + 1, val_loss_meter.avg))
print("Iter %d val_loss: %.4f" % (i + 1, val_loss_meter.avg))
# output scores
score, class_iou = running_metrics_val.get_scores()
for k, v in score.items():
print(k, v)
sys.stdout.flush()
logger.info('{}: {}'.format(k, v))
writer.add_scalar('val_metrics/{}'.format(k), v, i + 1)
for k, v in class_iou.items():
logger.info('{}: {}'.format(k, v))
writer.add_scalar('val_metrics/cls_{}'.format(k), v, i + 1)
torch.cuda.empty_cache()
# SAVE best model for segmentation
save_model = False
if score["Mean IoU : \t"] >= best_iou:
best_iou = score["Mean IoU : \t"]
save_model = True
if save_model:
# SAVE
savefilename = args.savemodel + 'model_best_iou'+ '.tar'
torch.save({
'iter': i,
'state_dict': model.state_dict(),
"best_iou": best_iou,
}, savefilename)
val_loss_meter.reset()
running_metrics_val.reset()
if (i + 1) == train_iters:
flag = False
break
i += 1
