def train(args):
def type_callback(event):
# print('event_type:{}'.format(event['event_type']))
if event['event_type'] == 'KeyPress':
event_key = event['key']
if event_key == 'Enter':
pass
# print('event_type:Enter')
elif event_key == 'Backspace':
pass
# print('event_type:Backspace')
elif event_key == 'Delete':
pass
# print('event_type:Delete')
elif len(event_key) == 1:
pass
# print('event_key:{}'.format(event['key']))
if event_key == 's':
import json
win = 'loss_iteration'
win_data = vis.get_window_data(win)
win_data_dict = json.loads(win_data)
win_data_content_dict = win_data_dict['content']
win_data_x = np.array(
win_data_content_dict['data'][0]['x'])
win_data_y = np.array(
win_data_content_dict['data'][0]['y'])
win_data_save_file = '/tmp/loss_iteration_{}.txt'.format(
init_time)
with open(win_data_save_file, 'wb') as f:
for item_x, item_y in zip(win_data_x, win_data_y):
f.write("{} {}\n".format(item_x, item_y))
done_time = str(int(time.time()))
vis.text(vis_text_usage + 'done at {}'.format(done_time),
win=callback_text_usage_window)
init_time = str(int(time.time()))
if args.vis:
# start visdom and close all window
vis = visdom.Visdom()
vis.close()
vis_text_usage = 'Operating in the text window<br>Press s to save data<br>'
callback_text_usage_window = vis.text(vis_text_usage)
vis.register_event_handler(type_callback, callback_text_usage_window)
class_weight = None
local_path = os.path.expanduser(args.dataset_path)
train_dst = None
val_dst = None
if args.dataset == 'SegmPred':
input_channel = 19
train_dst = segmpredLoader(local_path,
is_transform=True,
split='train')
val_dst = segmpredLoader(local_path, is_transform=True, split='val')
elif args.dataset == 'MovingMNIST':
# class_weight = [0.1, 0.5]
# class_weight = torch.tensor(class_weight)
train_dst = movingmnistLoader(local_path,
is_transform=True,
split='train')
val_dst = movingmnistLoader(local_path, is_transform=True, split='val')
elif args.dataset == 'FreeSpacePred':
input_channel = 1
train_dst = freespacepredLoader(local_path,
is_transform=True,
split='train')
val_dst = freespacepredLoader(local_path,
is_transform=True,
split='test')
else:
print('{} dataset does not implement'.format(args.dataset))
exit(0)
if args.cuda:
if class_weight is not None:
class_weight = class_weight.cuda()
print('class_weight:', class_weight)
train_loader = torch.utils.data.DataLoader(train_dst,
batch_size=args.batch_size,
shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dst,
batch_size=1,
shuffle=True)
start_epoch = 0
best_mIoU = 0
if args.resume_model != '':
model = torch.load(args.resume_model)
start_epoch_id1 = args.resume_model.rfind('_')
start_epoch_id2 = args.resume_model.rfind('.')
start_epoch = int(args.resume_model[start_epoch_id1 +
1:start_epoch_id2])
else:
# ---------------for testing SegmPred---------------
model = eval(args.structure)(n_classes=args.n_classes,
pretrained=args.init_vgg16,
input_shape=train_dst.input_shape,
input_channel=input_channel)
val_model = eval(args.structure)(n_classes=args.n_classes,
pretrained=args.init_vgg16,
input_shape=val_dst.input_shape,
input_channel=input_channel)
# ---------------for testing SegmPred---------------
if args.resume_model_state_dict != '':
try:
# from model save format get useful information, such as miou, epoch
miou_model_name_str = '_miou_'
class_model_name_str = '_class_'
miou_id1 = args.resume_model_state_dict.find(
miou_model_name_str) + len(miou_model_name_str)
miou_id2 = args.resume_model_state_dict.find(
class_model_name_str)
best_mIoU = float(
args.resume_model_state_dict[miou_id1:miou_id2])
start_epoch_id1 = args.resume_model_state_dict.rfind('_')
start_epoch_id2 = args.resume_model_state_dict.rfind('.')
start_epoch = int(
args.resume_model_state_dict[start_epoch_id1 +
1:start_epoch_id2])
pretrained_dict = torch.load(args.resume_model_state_dict,
map_location='cpu')
model.load_state_dict(pretrained_dict)
except KeyError:
print('missing resume_model_state_dict or wrong type')
if args.cuda:
model.cuda()
val_model.cuda()
print('start_epoch:', start_epoch)
print('best_mIoU:', best_mIoU)
if args.solver == 'SGD':
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
momentum=0.99,
weight_decay=5e-4)
elif args.solver == 'RMSprop':
optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
momentum=0.99,
weight_decay=5e-4)
elif args.solver == 'Adam':
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=5e-4)
else:
print('missing solver or not support')
exit(0)
# when observerd object dose not decrease scheduler will let the optimizer learing rate decrease
# scheduler = ReduceLROnPlateau(optimizer, 'min', patience=100, min_lr=1e-10, verbose=True)
if args.lr_policy == 'Constant':
scheduler = ConstantLR(optimizer)
elif args.lr_policy == 'Polynomial':
scheduler = PolynomialLR(
optimizer, max_iter=args.training_epoch,
power=0.9) # base lr=0.01 power=0.9 like PSPNet
elif args.lr_policy == 'MultiStep':
scheduler = MultiStepLR(
optimizer, milestones=[10, 50, 90],
gamma=0.1) # base lr=0.01 power=0.9 like PSPNet
# scheduler = StepLR(optimizer, step_size=1, gamma=0.1)
data_count = int(train_dst.__len__() * 1.0 / args.batch_size)
print('data_count:', data_count)
# iteration_step = 0
train_gts, train_preds = [], []
for epoch in range(start_epoch + 1, args.training_epoch, 1):
loss_epoch = 0
scheduler.step()
optimizer.zero_grad(
) # when train next time zero all grad, just acc the grad when the epoch training
for i, (imgs, labels) in enumerate(train_loader):
# if i==1:
# break
model.train()
# 最后的几张图片可能不到batch_size的数量，比如batch_size=4，可能只剩3张
imgs_batch = imgs.shape[0]
if imgs_batch != args.batch_size:
break
# iteration_step += 1
imgs = Variable(imgs)
labels = Variable(labels)
if args.cuda:
imgs = imgs.cuda()
labels = labels.cuda()
outputs = model(imgs)
# print('imgs.size:', imgs.size())
# print('labels.size:', labels.size())
# print('outputs.size:', outputs.size())
loss = cross_entropy2d(outputs, labels, weight=class_weight)
# add grad backward the avg loss
loss_grad_acc_avg = loss * 1.0 / args.grad_acc_steps
loss_grad_acc_avg.backward()
loss_np = loss.cpu().data.numpy()
loss_epoch += loss_np
if (i + 1) % args.grad_acc_steps == 0:
optimizer.step()
# 一次backward后如果不清零，梯度是累加的
optimizer.zero_grad()
# ------------------train metris-------------------------------
train_pred = outputs.cpu().data.max(1)[1].numpy()
train_gt = labels.cpu().data.numpy()
for train_gt_, train_pred_ in zip(train_gt, train_pred):
train_gts.append(train_gt_)
train_preds.append(train_pred_)
# ------------------train metris-------------------------------
if args.vis and i % 50 == 0:
pred_labels = outputs.cpu().data.max(1)[1].numpy()
label_color = train_dst.decode_segmap(
labels.cpu().data.numpy()[0]).transpose(2, 0, 1)
pred_label_color = train_dst.decode_segmap(
pred_labels[0]).transpose(2, 0, 1)
win = 'label_color'
vis.image(label_color,
win=win,
opts=dict(title='Gt', caption='Ground Truth'))
win = 'pred_label_color'
vis.image(pred_label_color,
win=win,
opts=dict(title='Pred', caption='Prediction'))
# 显示一个周期的loss曲线
if args.vis:
win = 'loss_iteration'
loss_np_expand = np.expand_dims(loss_np, axis=0)
win_res = vis.line(X=np.ones(1) * (i + data_count *
(epoch - 1) + 1),
Y=loss_np_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * (i + data_count * (epoch - 1) + 1),
Y=loss_np_expand,
win=win,
opts=dict(title=win,
xlabel='iteration',
ylabel='loss'))
# val result on val dataset and pick best to save
if args.val_interval > 0 and epoch % args.val_interval == 0:
print('----starting val----')
# model.eval()
val_model.load_state_dict(model.state_dict())
val_model.eval()
val_gts, val_preds = [], []
for val_i, (val_imgs, val_labels) in enumerate(val_loader):
# print(val_i)
val_imgs = Variable(val_imgs, volatile=True)
val_labels = Variable(val_labels, volatile=True)
if args.cuda:
val_imgs = val_imgs.cuda()
val_labels = val_labels.cuda()
# val_outputs = model(val_imgs)
val_outputs = val_model(val_imgs)
val_pred = val_outputs.cpu().data.max(1)[1].numpy()
val_gt = val_labels.cpu().data.numpy()
for val_gt_, val_pred_ in zip(val_gt, val_pred):
val_gts.append(val_gt_)
val_preds.append(val_pred_)
score, class_iou = scores(val_gts,
val_preds,
n_class=args.n_classes)
for k, v in score.items():
print(k, v)
if k == 'Mean IoU : \t':
v_iou = v
if v > best_mIoU:
best_mIoU = v_iou
torch.save(
model.state_dict(),
'{}_{}_miou_{}_class_{}_{}.pt'.format(
args.structure, args.dataset, best_mIoU,
args.n_classes, epoch))
# 显示校准周期的mIoU
if args.vis:
win = 'mIoU_epoch'
v_iou_expand = np.expand_dims(v_iou, axis=0)
win_res = vis.line(X=np.ones(1) * epoch *
args.val_interval,
Y=v_iou_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch * args.val_interval,
Y=v_iou_expand,
win=win,
opts=dict(title=win,
xlabel='epoch',
ylabel='mIoU'))
for class_i in range(args.n_classes):
print(class_i, class_iou[class_i])
print('----ending val----')
# 显示多个周期的loss曲线
loss_avg_epoch = loss_epoch / (data_count * 1.0)
# print(loss_avg_epoch)
if args.vis:
win = 'loss_epoch'
loss_avg_epoch_expand = np.expand_dims(loss_avg_epoch, axis=0)
win_res = vis.line(X=np.ones(1) * epoch,
Y=loss_avg_epoch_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch,
Y=loss_avg_epoch_expand,
win=win,
opts=dict(title=win, xlabel='epoch', ylabel='loss'))
if args.vis:
win = 'lr_epoch'
lr_epoch = np.array(scheduler.get_lr())
lr_epoch_expand = np.expand_dims(lr_epoch, axis=0)
win_res = vis.line(X=np.ones(1) * epoch,
Y=lr_epoch_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch,
Y=lr_epoch_expand,
win=win,
opts=dict(title=win, xlabel='epoch', ylabel='lr'))
# ------------------train metris-------------------------------
if args.vis:
score, class_iou = scores(train_gts,
train_preds,
n_class=args.n_classes)
for k, v in score.items():
print(k, v)
if k == 'Overall Acc : \t':
# 显示校准周期的mIoU
overall_acc = v
if args.vis:
win = 'acc_epoch'
overall_acc_expand = np.expand_dims(overall_acc,
axis=0)
win_res = vis.line(X=np.ones(1) * epoch,
Y=overall_acc_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch,
Y=overall_acc_expand,
win=win,
opts=dict(title=win,
xlabel='epoch',
ylabel='accuracy'))
# clear for new training metrics
train_gts, train_preds = [], []
# ------------------train metris-------------------------------
if args.save_model and epoch % args.save_epoch == 0:
torch.save(
model.state_dict(),
'{}_{}_class_{}_{}.pt'.format(args.structure, args.dataset,
args.n_classes, epoch))
def train(args):
def type_callback(event):
# print('event_type:{}'.format(event['event_type']))
if event['event_type'] == 'KeyPress':
event_key = event['key']
if event_key == 'Enter':
pass
# print('event_type:Enter')
elif event_key == 'Backspace':
pass
# print('event_type:Backspace')
elif event_key == 'Delete':
pass
# print('event_type:Delete')
elif len(event_key) == 1:
pass
# print('event_key:{}'.format(event['key']))
if event_key=='s':
import json
win = 'loss_iteration'
win_data = vis.get_window_data(win)
win_data_dict = json.loads(win_data)
win_data_content_dict = win_data_dict['content']
win_data_x = np.array(win_data_content_dict['data'][0]['x'])
win_data_y = np.array(win_data_content_dict['data'][0]['y'])
win_data_save_file = '/tmp/loss_iteration_{}.txt'.format(init_time)
with open(win_data_save_file, 'wb') as f:
for item_x, item_y in zip(win_data_x, win_data_y):
f.write("{} {}\n".format(item_x, item_y))
done_time = str(int(time.time()))
vis.text(vis_text_usage+'done at {}'.format(done_time), win=callback_text_usage_window)
init_time = str(int(time.time()))
if args.vis:
# start visdom and close all window
vis = visdom.Visdom()
vis.close()
vis_text_usage = 'Operating in the text window<br>Press s to save data<br>'
callback_text_usage_window = vis.text(vis_text_usage)
vis.register_event_handler(type_callback, callback_text_usage_window)
class_weight = None
local_path = os.path.expanduser(args.dataset_path)
train_dst = None
val_dst = None
if args.dataset == 'CamVid':
train_dst = camvidLoader(local_path, is_transform=True, is_augment=args.data_augment, split='train')
val_dst = camvidLoader(local_path, is_transform=True, is_augment=False, split='val')
trainannot_image_dir = os.path.expanduser(os.path.join(local_path, "trainannot"))
trainannot_image_files = [os.path.join(trainannot_image_dir, file) for file in os.listdir(trainannot_image_dir) if file.endswith('.png')]
if args.class_weighting=='MFB':
class_weight = median_frequency_balancing(trainannot_image_files, num_classes=12)
class_weight = torch.tensor(class_weight)
elif args.class_weighting=='ENET':
class_weight = ENet_weighing(trainannot_image_files, num_classes=12)
class_weight = torch.tensor(class_weight)
elif args.dataset == 'CityScapes':
train_dst = cityscapesLoader(local_path, is_transform=True, split='train')
val_dst = cityscapesLoader(local_path, is_transform=True, split='val')
else:
print('{} dataset does not implement'.format(args.dataset))
exit(0)
if args.cuda:
if class_weight is not None:
class_weight = class_weight.cuda()
print('class_weight:', class_weight)
train_loader = torch.utils.data.DataLoader(train_dst, batch_size=args.batch_size, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dst, batch_size=1, shuffle=True)
yolo_B = 2
yolo_C = 2
yolo_S = 7
yolo_out_tensor_shape = yolo_B * 5 + yolo_C
print('yolo_out_tensor_shape:', yolo_out_tensor_shape)
det_criterion = yoloLoss(yolo_S, yolo_B, yolo_C, 5, 0.5, args.cuda)
det_file_root = os.path.expanduser('~/Data/CamVid/train/')
det_train_dst = yoloDataset(root=det_file_root, list_file=['camvid_det.txt'], train=True, transform=[transforms.ToTensor()], yolo_out_tensor_shape=yolo_out_tensor_shape)
det_train_loader = torch.utils.data.DataLoader(det_train_dst, batch_size=1, shuffle=True, num_workers=4)
start_epoch = 0
best_mIoU = 0
if args.resume_model != '':
model = torch.load(args.resume_model)
start_epoch_id1 = args.resume_model.rfind('_')
start_epoch_id2 = args.resume_model.rfind('.')
start_epoch = int(args.resume_model[start_epoch_id1+1:start_epoch_id2])
else:
model = drnsegmt_a_18(pretrained=args.init_vgg16, n_classes=args.n_classes, det_tensor_num=yolo_out_tensor_shape)
if args.resume_model_state_dict != '':
try:
# from model save format get useful information, such as miou, epoch
miou_model_name_str = '_miou_'
class_model_name_str = '_class_'
miou_id1 = args.resume_model_state_dict.find(miou_model_name_str)+len(miou_model_name_str)
miou_id2 = args.resume_model_state_dict.find(class_model_name_str)
best_mIoU = float(args.resume_model_state_dict[miou_id1:miou_id2])
start_epoch_id1 = args.resume_model_state_dict.rfind('_')
start_epoch_id2 = args.resume_model_state_dict.rfind('.')
start_epoch = int(args.resume_model_state_dict[start_epoch_id1 + 1:start_epoch_id2])
pretrained_dict = torch.load(args.resume_model_state_dict, map_location='cpu')
model.load_state_dict(pretrained_dict)
except KeyError:
print('missing resume_model_state_dict or wrong type')
if args.cuda:
model.cuda()
print('start_epoch:', start_epoch)
print('best_mIoU:', best_mIoU)
optimizer = None
if args.solver == 'SGD':
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, momentum=0.99, weight_decay=5e-4)
elif args.solver == 'RMSprop':
optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, momentum=0.99, weight_decay=5e-4)
elif args.solver == 'Adam':
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=5e-4)
else:
print('missing solver or not support')
exit(0)
# when observerd object dose not decrease scheduler will let the optimizer learing rate decrease
# scheduler = ReduceLROnPlateau(optimizer, 'min', patience=100, min_lr=1e-10, verbose=True)
scheduler = None
if args.lr_policy == 'Constant':
scheduler = ConstantLR(optimizer)
elif args.lr_policy == 'Polynomial':
scheduler = PolynomialLR(optimizer, max_iter=args.training_epoch, power=0.9) # base lr=0.01 power=0.9 like PSPNet
# scheduler = StepLR(optimizer, step_size=1, gamma=0.1)
data_count = int(train_dst.__len__() * 1.0 / args.batch_size)
det_data_count = int(det_train_dst.__len__() * 1.0 / 1)
print('data_count:', data_count)
# iteration_step = 0
train_gts, train_preds = [], []
for epoch in range(start_epoch+1, args.training_epoch, 1):
loss_epoch = 0
scheduler.step()
# ----for object detection----
for det_i, (det_imgs, det_labels, _) in enumerate(det_train_loader):
model.train()
# print('det_imgs.shape:', det_imgs.shape)
# print('det_labels.shape:', det_labels.shape)
det_imgs = Variable(det_imgs)
det_labels = Variable(det_labels)
if args.cuda:
det_imgs = det_imgs.cuda()
det_labels = det_labels.cuda()
_, outputs_det = model(det_imgs)
# print('outpust_det:', outputs_det.shape)
det_loss = det_criterion(outputs_det, det_labels)
det_loss = 0.02 * det_loss # for balance with segment and detection
det_loss_np = det_loss.cpu().data.numpy()
optimizer.zero_grad()
det_loss.backward()
optimizer.step()
# 显示一个周期的loss曲线
if args.vis:
win = 'det_loss_iteration'
det_loss_np_expand = np.expand_dims(det_loss_np, axis=0)
win_res = vis.line(X=np.ones(1)*(det_i+det_data_count*(epoch-1)+1), Y=det_loss_np_expand, win=win, update='append')
if win_res != win:
vis.line(X=np.ones(1)*(det_i+det_data_count*(epoch-1)+1), Y=det_loss_np_expand, win=win, opts=dict(title=win, xlabel='iteration', ylabel='loss'))
# ----for object detection----
# ----for semantic segment----
for i, (imgs, labels) in enumerate(train_loader):
# if i==1:
# break
model.train()
# 最后的几张图片可能不到batch_size的数量，比如batch_size=4，可能只剩3张
imgs_batch = imgs.shape[0]
if imgs_batch != args.batch_size:
break
# iteration_step += 1
imgs = Variable(imgs)
labels = Variable(labels)
if args.cuda:
imgs = imgs.cuda()
labels = labels.cuda()
outputs_sem, _ = model(imgs)
# print('outputs_sem.shape:', outputs_sem.shape)
# 一次backward后如果不清零，梯度是累加的
optimizer.zero_grad()
# print('outputs.size:', outputs.size())
# print('labels.size:', labels.size())
loss = cross_entropy2d(outputs_sem, labels, weight=class_weight)
loss_np = loss.cpu().data.numpy()
loss_epoch += loss_np
loss.backward()
optimizer.step()
# ------------------train metris-------------------------------
train_pred = outputs_sem.cpu().data.max(1)[1].numpy()
train_gt = labels.cpu().data.numpy()
for train_gt_, train_pred_ in zip(train_gt, train_pred):
train_gts.append(train_gt_)
train_preds.append(train_pred_)
# ------------------train metris-------------------------------
if args.vis and i%50==0:
pred_labels = outputs_sem.cpu().data.max(1)[1].numpy()
label_color = train_dst.decode_segmap(labels.cpu().data.numpy()[0]).transpose(2, 0, 1)
pred_label_color = train_dst.decode_segmap(pred_labels[0]).transpose(2, 0, 1)
win = 'label_color'
vis.image(label_color, win=win, opts=dict(title='Gt', caption='Ground Truth'))
win = 'pred_label_color'
vis.image(pred_label_color, win=win, opts=dict(title='Pred', caption='Prediction'))
# 显示一个周期的loss曲线
if args.vis:
win = 'loss_iteration'
loss_np_expand = np.expand_dims(loss_np, axis=0)
win_res = vis.line(X=np.ones(1)*(i+data_count*(epoch-1)+1), Y=loss_np_expand, win=win, update='append')
if win_res != win:
vis.line(X=np.ones(1)*(i+data_count*(epoch-1)+1), Y=loss_np_expand, win=win, opts=dict(title=win, xlabel='iteration', ylabel='loss'))
# ----for semantic segment----
# val result on val dataset and pick best to save
if args.val_interval > 0 and epoch % args.val_interval == 0:
print('----starting val----')
model.eval()
val_gts, val_preds = [], []
for val_i, (val_imgs, val_labels) in enumerate(val_loader):
# print(val_i)
val_imgs = Variable(val_imgs)
val_labels = Variable(val_labels)
if args.cuda:
val_imgs = val_imgs.cuda()
val_labels = val_labels.cuda()
val_outputs_sem, _ = model(val_imgs)
val_pred = val_outputs_sem.cpu().data.max(1)[1].numpy()
val_gt = val_labels.cpu().data.numpy()
for val_gt_, val_pred_ in zip(val_gt, val_pred):
val_gts.append(val_gt_)
val_preds.append(val_pred_)
score, class_iou = scores(val_gts, val_preds, n_class=args.n_classes)
for k, v in score.items():
print(k, v)
if k == 'Mean IoU : \t':
v_iou = v
if v > best_mIoU:
best_mIoU = v_iou
torch.save(model.state_dict(), '{}_{}_miou_{}_class_{}_{}.pt'.format(args.structure, args.dataset, best_mIoU, args.n_classes, epoch))
# 显示校准周期的mIoU
if args.vis:
win = 'mIoU_epoch'
v_iou_expand = np.expand_dims(v_iou, axis=0)
win_res = vis.line(X=np.ones(1)*epoch*args.val_interval, Y=v_iou_expand, win=win, update='append')
if win_res != win:
vis.line(X=np.ones(1)*epoch*args.val_interval, Y=v_iou_expand, win=win, opts=dict(title=win, xlabel='epoch', ylabel='mIoU'))
# for class_i in range(args.n_classes):
# print(class_i, class_iou[class_i])
print('----ending val----')
# 显示多个周期的loss曲线
loss_avg_epoch = loss_epoch / (data_count * 1.0)
# print(loss_avg_epoch)
if args.vis:
win = 'loss_epoch'
loss_avg_epoch_expand = np.expand_dims(loss_avg_epoch, axis=0)
win_res = vis.line(X=np.ones(1)*epoch, Y=loss_avg_epoch_expand, win=win, update='append')
if win_res != win:
vis.line(X=np.ones(1)*epoch, Y=loss_avg_epoch_expand, win=win, opts=dict(title=win, xlabel='epoch', ylabel='loss'))
if args.vis:
win = 'lr_epoch'
lr_epoch = np.array(scheduler.get_lr())
lr_epoch_expand = np.expand_dims(lr_epoch, axis=0)
win_res = vis.line(X=np.ones(1)*epoch, Y=lr_epoch_expand, win=win, update='append')
if win_res != win:
vis.line(X=np.ones(1)*epoch, Y=lr_epoch_expand, win=win, opts=dict(title=win, xlabel='epoch', ylabel='lr'))
# ------------------train metris-------------------------------
if args.vis:
score, class_iou = scores(train_gts, train_preds, n_class=args.n_classes)
for k, v in score.items():
print(k, v)
if k == 'Overall Acc : \t':
# 显示校准周期的mIoU
overall_acc = v
if args.vis:
win = 'acc_epoch'
overall_acc_expand = np.expand_dims(overall_acc, axis=0)
win_res = vis.line(X=np.ones(1) * epoch, Y=overall_acc_expand, win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch, Y=overall_acc_expand, win=win,
opts=dict(title=win, xlabel='epoch', ylabel='accuracy'))
# clear for new training metrics
train_gts, train_preds = [], []
# ------------------train metris-------------------------------
if args.save_model and epoch%args.save_epoch==0:
torch.save(model.state_dict(), '{}_{}_class_{}_{}.pt'.format(args.structure, args.dataset, args.n_classes, epoch))
def train(args):
now = datetime.datetime.now()
now_str = '{}-{}-{} {}:{}:{}'.format(now.year, now.month, now.day,
now.hour, now.minute, now.second)
# print('now:', now)
# print('now_str:', now_str)
if args.vis:
# start visdom and close all window
vis = visdom.Visdom(env=now_str)
vis.close()
class_weight = None
local_path = os.path.expanduser(args.dataset_path)
train_dst = None
val_dst = None
if args.dataset == 'CamVid':
train_dst = camvidLoader(local_path,
is_transform=True,
is_augment=args.data_augment,
split='train')
val_dst = camvidLoader(local_path,
is_transform=True,
is_augment=False,
split='val')
trainannot_image_dir = os.path.expanduser(
os.path.join(local_path, "trainannot"))
trainannot_image_files = [
os.path.join(trainannot_image_dir, file)
for file in os.listdir(trainannot_image_dir)
if file.endswith('.png')
]
if args.class_weighting == 'MFB':
class_weight = median_frequency_balancing(trainannot_image_files,
num_classes=12)
class_weight = torch.tensor(class_weight)
elif args.class_weighting == 'ENET':
class_weight = ENet_weighing(trainannot_image_files,
num_classes=12)
class_weight = torch.tensor(class_weight)
elif args.dataset == 'CityScapes':
train_dst = cityscapesLoader(local_path,
is_transform=True,
split='train')
val_dst = cityscapesLoader(local_path, is_transform=True, split='val')
elif args.dataset == 'SegmPred':
train_dst = segmpredLoader(local_path,
is_transform=True,
split='train')
val_dst = segmpredLoader(local_path, is_transform=True, split='train')
elif args.dataset == 'MovingMNIST':
# class_weight = [0.1, 0.5]
# class_weight = torch.tensor(class_weight)
train_dst = movingmnistLoader(local_path,
is_transform=True,
split='train')
val_dst = movingmnistLoader(local_path, is_transform=True, split='val')
elif args.dataset == 'FreeSpace':
train_dst = freespaceLoader(local_path,
is_transform=True,
split='train')
val_dst = freespaceLoader(local_path, is_transform=True, split='val')
else:
print('{} dataset does not implement'.format(args.dataset))
exit(0)
if args.cuda:
if class_weight is not None:
class_weight = class_weight.cuda()
print('class_weight:', class_weight)
train_loader = torch.utils.data.DataLoader(train_dst,
batch_size=args.batch_size,
shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dst,
batch_size=1,
shuffle=True)
start_epoch = 0
best_mIoU = 0
if args.resume_model != '':
model = torch.load(args.resume_model)
start_epoch_id1 = args.resume_model.rfind('_')
start_epoch_id2 = args.resume_model.rfind('.')
start_epoch = int(args.resume_model[start_epoch_id1 +
1:start_epoch_id2])
else:
# model = eval(args.structure)(n_classes=args.n_classes, pretrained=args.init_vgg16)
try:
model = eval(args.structure)(n_classes=args.n_classes,
pretrained=args.init_vgg16)
except:
print('missing structure or not support')
exit(0)
# ---------------for testing SegmPred---------------
if args.dataset == 'MovingMNIST':
input_channel = 1 * 9
elif args.dataset == 'SegmPred':
input_channel = 19 * 4
if args.structure == 'drnseg_a_18':
model = drnseg_a_18(n_classes=args.n_classes,
pretrained=args.init_vgg16,
input_channel=input_channel)
# ---------------for testing SegmPred---------------
if args.resume_model_state_dict != '':
try:
# from model save format get useful information, such as miou, epoch
miou_model_name_str = '_miou_'
class_model_name_str = '_class_'
miou_id1 = args.resume_model_state_dict.find(
miou_model_name_str) + len(miou_model_name_str)
miou_id2 = args.resume_model_state_dict.find(
class_model_name_str)
best_mIoU = float(
args.resume_model_state_dict[miou_id1:miou_id2])
start_epoch_id1 = args.resume_model_state_dict.rfind('_')
start_epoch_id2 = args.resume_model_state_dict.rfind('.')
start_epoch = int(
args.resume_model_state_dict[start_epoch_id1 +
1:start_epoch_id2])
pretrained_dict = torch.load(args.resume_model_state_dict,
map_location='cpu')
model.load_state_dict(pretrained_dict)
except KeyError:
print('missing resume_model_state_dict or wrong type')
if args.cuda:
model.cuda()
print('start_epoch:', start_epoch)
print('best_mIoU:', best_mIoU)
if args.solver == 'SGD':
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
momentum=0.99,
weight_decay=5e-4)
elif args.solver == 'RMSprop':
optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
momentum=0.99,
weight_decay=5e-4)
elif args.solver == 'Adam':
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=5e-4)
else:
print('missing solver or not support')
exit(0)
# when observerd object dose not decrease scheduler will let the optimizer learing rate decrease
# scheduler = ReduceLROnPlateau(optimizer, 'min', patience=100, min_lr=1e-10, verbose=True)
if args.lr_policy == 'Constant':
scheduler = ConstantLR(optimizer)
elif args.lr_policy == 'Polynomial':
scheduler = PolynomialLR(
optimizer, max_iter=args.training_epoch,
power=0.9) # base lr=0.01 power=0.9 like PSPNet
elif args.lr_policy == 'MultiStep':
scheduler = MultiStepLR(
optimizer, milestones=[10, 50, 90],
gamma=0.1) # base lr=0.01 power=0.9 like PSPNet
# scheduler = StepLR(optimizer, step_size=1, gamma=0.1)
data_count = int(train_dst.__len__() * 1.0 / args.batch_size)
print('data_count:', data_count)
# iteration_step = 0
train_gts, train_preds = [], []
for epoch in range(start_epoch + 1, args.training_epoch, 1):
loss_epoch = 0
scheduler.step()
optimizer.zero_grad(
) # when train next time zero all grad, just acc the grad when the epoch training
for i, (imgs, labels) in enumerate(train_loader):
# if i==1:
# break
model.train()
# 最后的几张图片可能不到batch_size的数量，比如batch_size=4，可能只剩3张
imgs_batch = imgs.shape[0]
if imgs_batch != args.batch_size:
break
# iteration_step += 1
imgs = Variable(imgs)
labels = Variable(labels)
if args.cuda:
imgs = imgs.cuda()
labels = labels.cuda()
outputs = model(imgs)
# print('imgs.size:', imgs.size())
# print('labels.size:', labels.size())
# print('outputs.size:', outputs.size())
loss = cross_entropy2d(outputs, labels, weight=class_weight)
# add grad backward the avg loss
loss_grad_acc_avg = loss * 1.0 / args.grad_acc_steps
loss_grad_acc_avg.backward()
loss_np = loss.cpu().data.numpy()
loss_epoch += loss_np
if (i + 1) % args.grad_acc_steps == 0:
optimizer.step()
# 一次backward后如果不清零，梯度是累加的
optimizer.zero_grad()
# ------------------train metris-------------------------------
train_pred = outputs.cpu().data.max(1)[1].numpy()
train_gt = labels.cpu().data.numpy()
for train_gt_, train_pred_ in zip(train_gt, train_pred):
train_gts.append(train_gt_)
train_preds.append(train_pred_)
# ------------------train metris-------------------------------
if args.vis and i % 50 == 0:
pred_labels = outputs.cpu().data.max(1)[1].numpy()
label_color = train_dst.decode_segmap(
labels.cpu().data.numpy()[0]).transpose(2, 0, 1)
pred_label_color = train_dst.decode_segmap(
pred_labels[0]).transpose(2, 0, 1)
win = 'label_color'
vis.image(label_color,
win=win,
opts=dict(title='Gt', caption='Ground Truth'))
win = 'pred_label_color'
vis.image(pred_label_color,
win=win,
opts=dict(title='Pred', caption='Prediction'))
# 显示一个周期的loss曲线
if args.vis:
win = 'loss_iteration'
loss_np_expand = np.expand_dims(loss_np, axis=0)
win_res = vis.line(X=np.ones(1) * (i + data_count *
(epoch - 1) + 1),
Y=loss_np_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * (i + data_count * (epoch - 1) + 1),
Y=loss_np_expand,
win=win,
opts=dict(title=win,
xlabel='iteration',
ylabel='loss'))
# val result on val dataset and pick best to save
if args.val_interval > 0 and epoch % args.val_interval == 0:
print('----starting val----')
model.eval()
val_gts, val_preds = [], []
for val_i, (val_imgs, val_labels) in enumerate(val_loader):
# print(val_i)
val_imgs = Variable(val_imgs, volatile=True)
val_labels = Variable(val_labels, volatile=True)
if args.cuda:
val_imgs = val_imgs.cuda()
val_labels = val_labels.cuda()
val_outputs = model(val_imgs)
val_pred = val_outputs.cpu().data.max(1)[1].numpy()
val_gt = val_labels.cpu().data.numpy()
for val_gt_, val_pred_ in zip(val_gt, val_pred):
val_gts.append(val_gt_)
val_preds.append(val_pred_)
score, class_iou = scores(val_gts,
val_preds,
n_class=args.n_classes)
for k, v in score.items():
print(k, v)
if k == 'Mean IoU : \t':
v_iou = v
if v > best_mIoU:
best_mIoU = v_iou
torch.save(
model.state_dict(),
'{}_{}_miou_{}_class_{}_{}.pt'.format(
args.structure, args.dataset, best_mIoU,
args.n_classes, epoch))
# 显示校准周期的mIoU
if args.vis:
win = 'mIoU_epoch'
v_iou_expand = np.expand_dims(v_iou, axis=0)
win_res = vis.line(X=np.ones(1) * epoch *
args.val_interval,
Y=v_iou_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch * args.val_interval,
Y=v_iou_expand,
win=win,
opts=dict(title=win,
xlabel='epoch',
ylabel='mIoU'))
for class_i in range(args.n_classes):
print(class_i, class_iou[class_i])
print('----ending val----')
# 显示多个周期的loss曲线
loss_avg_epoch = loss_epoch / (data_count * 1.0)
# print(loss_avg_epoch)
if args.vis:
win = 'loss_epoch'
loss_avg_epoch_expand = np.expand_dims(loss_avg_epoch, axis=0)
win_res = vis.line(X=np.ones(1) * epoch,
Y=loss_avg_epoch_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch,
Y=loss_avg_epoch_expand,
win=win,
opts=dict(title=win, xlabel='epoch', ylabel='loss'))
if args.vis:
win = 'lr_epoch'
lr_epoch = np.array(scheduler.get_lr())
lr_epoch_expand = np.expand_dims(lr_epoch, axis=0)
win_res = vis.line(X=np.ones(1) * epoch,
Y=lr_epoch_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch,
Y=lr_epoch_expand,
win=win,
opts=dict(title=win, xlabel='epoch', ylabel='lr'))
# ------------------train metris-------------------------------
if args.vis:
score, class_iou = scores(train_gts,
train_preds,
n_class=args.n_classes)
for k, v in score.items():
print(k, v)
if k == 'Overall Acc : \t':
# 显示校准周期的mIoU
overall_acc = v
if args.vis:
win = 'acc_epoch'
overall_acc_expand = np.expand_dims(overall_acc,
axis=0)
win_res = vis.line(X=np.ones(1) * epoch,
Y=overall_acc_expand,
win=win,
update='append')
if win_res != win:
vis.line(X=np.ones(1) * epoch,
Y=overall_acc_expand,
win=win,
opts=dict(title=win,
xlabel='epoch',
ylabel='accuracy'))
# clear for new training metrics
train_gts, train_preds = [], []
# ------------------train metris-------------------------------
if args.save_model and epoch % args.save_epoch == 0:
torch.save(
model.state_dict(),
'{}_{}_class_{}_{}.pt'.format(args.structure, args.dataset,
args.n_classes, epoch))