def main():
net = PSPNet(num_classes=voc.num_classes).cuda()
print('load model ' + args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, args['exp_name'], args['snapshot'])))
net.eval()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
val_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
test_set = voc.VOC('test', transform=val_input_transform)
test_loader = DataLoader(test_set, batch_size=1, num_workers=8, shuffle=False)
check_mkdir(os.path.join(ckpt_path, args['exp_name'], 'test'))
for vi, data in enumerate(test_loader):
img_name, img = data
img_name = img_name[0]
img = Variable(img, volatile=True).cuda()
output = net(img)
prediction = output.data.max(1)[1].squeeze_(1).squeeze_(0).cpu().numpy()
prediction = voc.colorize_mask(prediction)
prediction.save(os.path.join(ckpt_path, args['exp_name'], 'test', img_name + '.png'))
print('%d / %d' % (vi + 1, len(test_loader)))
def validate(val_loader, net, criterion, optimizer, epoch, iter_num, train_args, visualize):
# the following code is written assuming that batch size is 1
net.eval()
val_loss = AverageMeter()
gts_all = np.zeros((len(val_loader), args['longer_size'] / 2, args['longer_size']), dtype=int)
predictions_all = np.zeros((len(val_loader), args['longer_size'] / 2, args['longer_size']), dtype=int)
for vi, data in enumerate(val_loader):
input, gt, slices_info = data
assert len(input.size()) == 5 and len(gt.size()) == 4 and len(slices_info.size()) == 3
input.transpose_(0, 1)
gt.transpose_(0, 1)
slices_info.squeeze_(0)
assert input.size()[3:] == gt.size()[2:]
count = torch.zeros(args['longer_size'] / 2, args['longer_size']).cuda()
output = torch.zeros(cityscapes.num_classes, args['longer_size'] / 2, args['longer_size']).cuda()
slice_batch_pixel_size = input.size(1) * input.size(3) * input.size(4)
for input_slice, gt_slice, info in zip(input, gt, slices_info):
input_slice = Variable(input_slice).cuda()
gt_slice = Variable(gt_slice).cuda()
output_slice = net(input_slice)
assert output_slice.size()[2:] == gt_slice.size()[1:]
assert output_slice.size()[1] == cityscapes.num_classes
output[:, info[0]: info[1], info[2]: info[3]] += output_slice[0, :, :info[4], :info[5]].data
gts_all[vi, info[0]: info[1], info[2]: info[3]] += gt_slice[0, :info[4], :info[5]].data.cpu().numpy()
count[info[0]: info[1], info[2]: info[3]] += 1
val_loss.update(criterion(output_slice, gt_slice).data[0], slice_batch_pixel_size)
output /= count
gts_all[vi, :, :] /= count.cpu().numpy().astype(int)
predictions_all[vi, :, :] = output.max(0)[1].squeeze_(0).cpu().numpy()
print('validating: %d / %d' % (vi + 1, len(val_loader)))
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all, cityscapes.num_classes)
if val_loss.avg < train_args['best_record']['val_loss']:
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['iter'] = iter_num
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
snapshot_name = 'epoch_%d_iter_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.10f' % (
epoch, iter_num, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, optimizer.param_groups[1]['lr'])
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, '%d_%d' % (epoch, iter_num))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(zip(gts_all, predictions_all)):
gt_pil = cityscapes.colorize_mask(data[0])
predictions_pil = cityscapes.colorize_mask(data[1])
if train_args['val_save_to_img_file']:
predictions_pil.save(os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=2, padding=5)
writer.add_image(snapshot_name, val_visual)
print('-----------------------------------------------------------------------------------------------------------')
print('[epoch %d], [iter %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' % (
epoch, iter_num, val_loss.avg, acc, acc_cls, mean_iu, fwavacc))
print('best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d], '
'[iter %d]' % (train_args['best_record']['val_loss'], train_args['best_record']['acc'],
train_args['best_record']['acc_cls'], train_args['best_record']['mean_iu'],
train_args['best_record']['fwavacc'], train_args['best_record']['epoch'],
train_args['best_record']['iter']))
print('-----------------------------------------------------------------------------------------------------------')
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
net.train()
return val_loss.avg
def main():
print('===> Loading datasets')
train_set = get_training_set(opt.data_dir, opt.train_dir, opt.patch_size,
opt.sr_patch_size, opt.sr_upscale_factor,
opt.num_classes, opt.sr_data_augmentation)
if opt.val_dir != None:
val_set = get_eval_set(opt.data_dir, opt.val_dir,
opt.sr_upscale_factor, opt.num_classes)
train_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batch_size)
val_loader = DataLoader(dataset=val_set,
num_workers=opt.threads,
batch_size=1)
else:
# Creating data indices for training and validation splits:
validation_split = .2
dataset_size = len(train_set)
indices = list(range(dataset_size))
split = int(np.floor(validation_split * dataset_size))
np.random.seed(opt.seed)
np.random.shuffle(indices)
train_indices, val_indices = indices[split:], indices[:split]
train_sampler = SubsetRandomSampler(train_indices)
val_sampler = SubsetRandomSampler(val_indices)
train_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batch_size,
sampler=train_sampler)
val_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=1,
sampler=val_sampler)
print('Building SR model ', opt.sr_model_name)
if opt.sr_model_name == 'DBPN':
sr_model = DBPN(num_channels=3,
base_filter=64,
feat=256,
num_stages=7,
scale_factor=opt.sr_upscale_factor)
sr_model = torch.nn.DataParallel(sr_model, device_ids=gpus_list)
if opt.sr_pretrained:
model_name = os.path.join(opt.save_folder +
opt.sr_pretrained_model)
print(model_name)
sr_model.load_state_dict(
torch.load(model_name,
map_location=lambda storage, loc: storage))
print('Pre-trained SR model is loaded.')
else:
sys.exit('Invalid SR network')
print('Building SemSeg model', opt.seg_model_name)
if opt.seg_model_name == 'segnet':
seg_model = segnet(num_classes=opt.num_classes, in_channels=3)
if not opt.seg_pretrained:
seg_model.init_vgg16_params()
print('segnet params initialized')
seg_model = torch.nn.DataParallel(seg_model, device_ids=gpus_list)
if opt.seg_pretrained:
model_name = os.path.join(opt.save_folder +
opt.seg_pretrained_model)
print(model_name)
seg_model.load_state_dict(torch.load(model_name))
print('Pre-trained SemSeg model is loaded.')
seg_model = torch.nn.DataParallel(seg_model, device_ids=gpus_list)
sr_criterion = nn.L1Loss()
psnr_criterion = nn.MSELoss()
if cuda:
sr_model = sr_model.cuda(gpus_list[0])
seg_model = seg_model.cuda(gpus_list[0])
sr_criterion = sr_criterion.cuda(gpus_list[0])
psnr_criterion = psnr_criterion.cuda(gpus_list[0])
if 'grss' in opt.data_dir:
seg_criterion = CrossEntropyLoss2d(ignore_index=-1).cuda()
else:
seg_criterion = CrossEntropyLoss2d().cuda()
sr_optimizer = optim.Adam(sr_model.parameters(),
lr=opt.sr_lr,
betas=(0.9, 0.999),
eps=1e-8)
seg_optimizer = optim.Adam(seg_model.parameters(),
lr=opt.seg_lr,
weight_decay=opt.seg_weight_decay,
betas=(opt.seg_momentum, 0.99))
scheduler = ReduceLROnPlateau(seg_optimizer,
'min',
factor=0.5,
patience=opt.seg_lr_patience,
min_lr=2.5e-5,
verbose=True)
check_mkdir(os.path.join('outputs', exp_name))
check_mkdir(os.path.join('outputs', exp_name, 'segmentation'))
check_mkdir(os.path.join('outputs', exp_name, 'super-resolution'))
check_mkdir(os.path.join(opt.save_folder, exp_name))
#best_iou = 0
best_iou = val_results = validate(0, val_loader, sr_model, seg_model,
sr_criterion, psnr_criterion,
seg_criterion, sr_optimizer,
seg_optimizer)
#sys.exit()
#best_epoch = -1
best_epoch = 0
best_model = (sr_model, seg_model)
since_last_best = 0
for epoch in range(opt.start_iter, opt.epoch_num + 1):
train(epoch, train_loader, sr_model, seg_model, sr_criterion,
psnr_criterion, seg_criterion, sr_optimizer, seg_optimizer)
val_results = validate(epoch, val_loader, sr_model, seg_model,
sr_criterion, psnr_criterion, seg_criterion,
sr_optimizer, seg_optimizer)
if val_results > best_iou:
best_iou = val_results
best_epoch = epoch
print('New best iou ', best_iou)
best_model = (copy.deepcopy(sr_model), copy.deepcopy(seg_model))
since_last_best = 0
checkpoint(epoch, sr_model, seg_model, 'tmp_best')
else:
print('Best iou epoch: ', best_epoch, ':', best_iou)
scheduler.step(val_results)
if (epoch) % (opt.epoch_num / 2) == 0:
for param_group in sr_optimizer.param_groups:
param_group['lr'] /= 10.0
print('SR Learning rate decay: lr={}'.format(
sr_optimizer.param_groups[0]['lr']))
if (epoch) % (opt.snapshots) == 0:
checkpoint(epoch, sr_model, seg_model)
#since_last_best += 1
#if since_last_best == 20:
# checkpoint(epoch, best_model[0], best_model[1], 'tmp_best')
print('Saving final best model')
checkpoint(epoch, best_model[0], best_model[1], 'best')
help='Batch size')
parser.add_argument('--lr', nargs='?', type=float, default=5e-4,
help='Learning Rate')
parser.add_argument('--weight_decay', nargs='?', type=float, default=5e-4,
help='Weight decay')
parser.add_argument('--momentum', nargs='?', type=float, default=0.9,
help='momentum')
parser.add_argument('--load_param', nargs='?', type=str, default=None,
help='Path to previous saved parameters to restart from')
parser.add_argument('--load_optim', nargs='?', type=str, default=None,
help='Path to previous saved optimizer to restart from')
args = parser.parse_args()
ckpt_path = '../../net_data/camvid'
exp_name = 'camvid_linknet_cat'
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(args) + '\n\n')
writer = SummaryWriter(os.path.join(ckpt_path, 'exp', exp_name))
main(args)
from os.path import join
from utils import check_mkdir
caffe_mode = "GPU" # GPU|CPU
caffe_device = 0
caffe_root = "/home/avp329/CV/Yelp/caffe/"
data_root = "/scratch/avp329/YelpData/"
check_mkdir(data_root + "feature_set")
check_mkdir("submission")
# Models: (model_weights, model_prototxt, mean_image, image_feature_file, biz_feature_file)
models = ((caffe_root + "models/bvlc_reference_caffenet.caffemodel",
caffe_root + "models/bvlc_reference_caffenet.prototxt",
caffe_root + "models/ilsvrc_2012_mean.npy",
data_root + "feature_set/bvlc_reference_{}_image_features.h5",
data_root + "feature_set/bvlc_reference_{}_biz_features_{}.pkl"),
(caffe_root + "models/places205CNN_iter_300000.caffemodel",
caffe_root + "models/places205CNN.prototxt",
caffe_root + "models/places205CNN_mean.npy",
data_root + "feature_set/places205CNN_{}_image_features.h5",
data_root + "feature_set/places205CNN_{}_biz_features_{}.pkl"),
(caffe_root + "models/hybridCNN_iter_700000.caffemodel",
caffe_root + "models/places205CNN.prototxt",
caffe_root + "models/hybridCNN_mean.npy",
data_root + "feature_set/hybridCNN_{}_image_features.h5",
data_root + "feature_set/hybridCNN_{}_biz_features_{}.pkl"),
(caffe_root + "models/bvlc_alexnet.caffemodel",
caffe_root + "models/bvlc_alexnet.prototxt",
caffe_root + "models/ilsvrc_2012_mean.npy",
for attr_name, space_name, regularizer, end_point in zip(
attr_name_list, space_name_list, regularizer_list, end_point_list):
# skip for tested attribute
if attr_name == 'smile' or attr_name == 'yaw' or attr_name == 'happiness':
continue
svm_path = os.path.join(svm_root, space_name, regularizer,
'n_{}.pt'.format(attr_name))
svm_data = torch.load(svm_path)
w = svm_data['w']
b = svm_data['b']
# create the saving directory
save_img_dir = os.path.join(edit_img_dir, attr_name)
check_mkdir(save_img_dir)
for latent_name in latent_list:
latent_path = os.path.join(input_latent_dir, latent_name)
latent_data = torch.load(latent_path)
latent_code = latent_data[space_name]
core_tensor = latent_data['c']
# generate the edited attribute and latent
latent_code_list = latent_interpolate(latent_code,
w,
b,
latent_mode=space_name,
end_point=end_point,
steps=steps)
img_list = []
def main(train_args):
backbone = ResNet()
backbone.load_state_dict(torch.load(
'./weight/resnet34-333f7ec4.pth'), strict=False)
net = Decoder34(num_classes=13, backbone=backbone).cuda()
D = discriminator(input_channels=16).cuda()
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {
'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0}
else:
print('training resumes from ' + train_args['snapshot'])
net.load_state_dict(torch.load(os.path.join(
ckpt_path, exp_name, train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])}
net.train()
D.train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage(),
])
visualize = standard_transforms.Compose([
standard_transforms.Scale(400),
standard_transforms.CenterCrop(400),
standard_transforms.ToTensor()
])
train_set = wp.Wp('train', transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=4,
num_workers=4, shuffle=True)
# val_set = wp.Wp('val', transform=input_transform,
# target_transform=target_transform)
# XR：所以这里本来就不能用到val？这里为什么不用一个val的数据集呢？
val_loader = DataLoader(train_set, batch_size=1,
num_workers=4, shuffle=False)
criterion = DiceLoss().cuda()
criterion_D = nn.BCELoss().cuda()
optimizer_AE = optim.Adam([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * train_args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': train_args['lr'], 'weight_decay': train_args['weight_decay']}
], betas=(train_args['momentum'], 0.999))
optimizer_D = optim.Adam([
{'params': [param for name, param in D.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * train_args['lr']},
{'params': [param for name, param in D.named_parameters() if name[-4:] != 'bias'],
'lr': train_args['lr'], 'weight_decay': train_args['weight_decay']}
], betas=(train_args['momentum'], 0.999))
if len(train_args['snapshot']) > 0:
optimizer_AE.load_state_dict(torch.load(os.path.join(
ckpt_path, exp_name, 'opt_' + train_args['snapshot'])))
optimizer_AE.param_groups[0]['lr'] = 2 * train_args['lr']
optimizer_AE.param_groups[1]['lr'] = train_args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) +
'.txt'), 'w').write(str(train_args) + '\n\n')
scheduler = ReduceLROnPlateau(
optimizer_AE, 'min', patience=train_args['lr_patience'], min_lr=1e-10, verbose=True)
for epoch in range(curr_epoch, train_args['epoch_num'] + 1):
train(train_loader, net, D, criterion, criterion_D, optimizer_AE,
optimizer_D, epoch, train_args)
val_loss = validate(val_loader, net, criterion, optimizer_AE,
epoch, train_args, restore_transform, visualize)
scheduler.step(val_loss)
def validate(val_loader, net, criterion, optimizer, epoch, train_args, visualize):
# the following code is written assuming that batch size is 1
net.eval()
val_loss = AverageMeter()
gts_all = np.zeros((len(val_loader), args['shorter_size'], 2 * args['shorter_size']), dtype=int)
predictions_all = np.zeros((len(val_loader), args['shorter_size'], 2 * args['shorter_size']), dtype=int)
for vi, data in enumerate(val_loader):
input, gt, slices_info = data
assert len(input.size()) == 5 and len(gt.size()) == 4 and len(slices_info.size()) == 3
input.transpose_(0, 1)
gt.transpose_(0, 1)
slices_info.squeeze_(0)
assert input.size()[3:] == gt.size()[2:]
count = torch.zeros(args['shorter_size'], 2 * args['shorter_size']).cuda()
output = torch.zeros(voc.num_classes, args['shorter_size'], 2 * args['shorter_size']).cuda()
slice_batch_pixel_size = input.size(1) * input.size(3) * input.size(4)
for input_slice, gt_slice, info in zip(input, gt, slices_info):
input_slice = Variable(input_slice).cuda()
gt_slice = Variable(gt_slice).cuda()
output_slice = net(input_slice)
assert output_slice.size()[2:] == gt_slice.size()[1:]
assert output_slice.size()[1] == voc.num_classes
output[:, info[0]: info[1], info[2]: info[3]] += output_slice[0, :, :info[4], :info[5]].data
gts_all[vi, info[0]: info[1], info[2]: info[3]] += gt_slice[0, :info[4], :info[5]].data.cpu().numpy()
count[info[0]: info[1], info[2]: info[3]] += 1
val_loss.update(criterion(output_slice, gt_slice).data[0], slice_batch_pixel_size)
output /= count
gts_all[vi, :, :] /= count.cpu().numpy().astype(int)
predictions_all[vi, :, :] = output.max(0)[1].squeeze_(0).cpu().numpy()
print('validating: %d / %d' % (vi + 1, len(val_loader)))
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all, voc.num_classes)
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
snapshot_name = 'epoch_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.10f' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, optimizer.param_groups[1]['lr'])
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, str(epoch))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(zip(gts_all, predictions_all)):
gt_pil = voc.colorize_mask(data[0])
predictions_pil = voc.colorize_mask(data[1])
if train_args['val_save_to_img_file']:
predictions_pil.save(os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=2, padding=5)
writer.add_image(snapshot_name, val_visual)
print('-----------------------------------------------------------------------------------------------------------')
print('[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc))
print('best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d]' % (
train_args['best_record']['val_loss'], train_args['best_record']['acc'], train_args['best_record']['acc_cls'],
train_args['best_record']['mean_iu'], train_args['best_record']['fwavacc'], train_args['best_record']['epoch']))
print('-----------------------------------------------------------------------------------------------------------')
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
net.train()
return val_loss.avg
def validate(val_loader,
net,
device,
criterion,
optimizer,
epoch,
train_args,
restore,
visualize,
finetuning=False):
net.eval()
val_loss = AverageMeter()
inputs_all, gts_all, predictions_all = [], [], []
with torch.no_grad():
for vi, data in enumerate(val_loader):
inputs, gts = data
N = inputs.size(0)
inputs, gts = inputs.to(device), gts.to(device)
outputs = net(inputs)
predictions = outputs.data.max(1)[1].squeeze_(1).cpu().numpy()
val_loss.update(criterion(outputs, gts).data / N, N)
for i in inputs:
if random.random() > train_args['val_img_sample_rate']:
inputs_all.append(None)
else:
inputs_all.append(i.data.cpu())
gts_all.append(gts.data.cpu().numpy())
predictions_all.append(predictions)
gts_all = np.concatenate(gts_all)
predictions_all = np.concatenate(predictions_all)
acc, acc_cls, mean_iu, _ = evaluate(predictions_all, gts_all,
cityscapes.num_classes)
if mean_iu > train_args['best_record']['mean_iu']:
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
snapshot_name = 'epoch_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu)
if finetuning == False:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, str(epoch))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(
zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = restore(data[0])
gt_pil = cityscapes.colorize_mask(data[1])
predictions_pil = cityscapes.colorize_mask(data[2])
if train_args['val_save_to_img_file']:
input_pil.save(
os.path.join(to_save_dir, '%d_input.png' % idx))
predictions_pil.save(
os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([
visualize(input_pil.convert('RGB')),
visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))
])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=3, padding=5)
writer.add_image(snapshot_name, val_visual)
else:
torch.save(
net.state_dict(),
os.path.join(ft_ckpt_path, exp_name, snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ft_ckpt_path, exp_name, str(epoch))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(
zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = restore(data[0])
gt_pil = cityscapes.colorize_mask(data[1])
predictions_pil = cityscapes.colorize_mask(data[2])
if train_args['val_save_to_img_file']:
input_pil.save(
os.path.join(to_save_dir, '%d_input.png' % idx))
predictions_pil.save(
os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([
visualize(input_pil.convert('RGB')),
visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))
])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=3, padding=5)
writer.add_image(snapshot_name, val_visual)
print(
'-----------------------------------------------------------------------------------------------------------'
)
print(
'[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f]'
% (epoch, val_loss.avg, acc, acc_cls, mean_iu))
print(
'best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [epoch %d]'
% (train_args['best_record']['val_loss'],
train_args['best_record']['acc'],
train_args['best_record']['acc_cls'],
train_args['best_record']['mean_iu'],
train_args['best_record']['epoch']))
print(
'-----------------------------------------------------------------------------------------------------------'
)
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
net.train()
return val_loss.avg
def main(train_args):
check_mkdir(os.path.join(train_args['ckpt_path'], args['exp']))
check_mkdir(
os.path.join(train_args['ckpt_path'], args['exp'],
train_args['exp_name']))
model = DeepLabV3('1')
# print(model)
device = torch.device("cuda")
num_gpu = list(range(torch.cuda.device_count()))
"""###############------use gpu--------###############"""
if args['use_gpu']:
ts = time.time()
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(0))
model = nn.DataParallel(model, device_ids=num_gpu)
model = model.to(device)
print("Finish cuda loading ,time elapsed {}", format(time.time() - ts))
else:
print("please check your gpu device,start training on cpu")
"""###############-------中间开始训练--------###############"""
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
# model.apply(weights_init)
else:
print("train resume from " + train_args['snapshot'])
state_dict = torch.load(
os.path.join(train_args['ckpt_path'], args['exp'],
train_args['exp_name'], train_args['snapshot']))
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:]
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
# model.load_state_dict(
# torch.load(os.path.join(train_args['ckpt_path'],args['exp'],train_args['exp_name'], train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]),
'fwavacc': float(split_snapshot[11])
}
model.train()
mean_std = ([0.485, 0.456, 0.406, 0.450], [0.229, 0.224, 0.225, 0.225])
"""#################---数据增强和数据变换等操作------########"""
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
]) ##Nomorlized
target_transform = extended_transforms.MaskToTensor() # target to tensor
joint_transform = joint_transforms.Compose([
joint_transforms.RandomHorizontallyFlip(),
joint_transforms.RandomCrop((256, 256), padding=0),
joint_transforms.Rotate(degree=90)
]) ###data_augment
restore = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
extended_transforms.channel_4_to_channel_3(4, 3), ##默认3通道如果四通道会转成三通道
standard_transforms.ToPILImage(),
]) # DeNomorlized，出来是pil图片了
visualize = standard_transforms.Compose([
standard_transforms.Resize(256),
standard_transforms.CenterCrop(256), ##中心裁剪，此处可以删除
standard_transforms.ToTensor()
]) # resize 大小之后转tensor
"""#################---数据加载------########"""
train_set = yaogan(mode='train',
cls=train_args['training_cls'],
joint_transform=None,
input_transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=train_args['batch_size'],
num_workers=train_args['num_works'],
shuffle=True)
val_set = yaogan(mode='val',
cls=train_args['training_cls'],
input_transform=input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=1,
num_workers=train_args['num_works'],
shuffle=False)
# test_set=yaogan(mode='test',cls=train_args['training_cls'],joint_transform=None,
# input_transform=input_transform,target_transform=None)
# test_loader=DataLoader(test_set,batch_size=1,
# num_workers=train_args['num_works'], shuffle=False)
optimizer = optim.Adadelta(model.parameters(), lr=train_args['lr'])
##define a weighted loss (0weight for 0 label)
# weight=[0.09287939 ,0.02091968 ,0.02453979, 0.25752962 ,0.33731845, 1.,
# 0.09518322, 0.52794035 ,0.24298112 ,0.02657369, 0.15057124 ,0.36864611,
# 0.25835161,0.16672758 ,0.40728756 ,0.00751281]
"""###############-------训练数据权重--------###############"""
if train_args['weight'] is not None:
weight = [0.1, 1.]
weight = torch.Tensor(weight)
else:
weight = None
criterion = nn.CrossEntropyLoss(weight=weight,
reduction='elementwise_mean',
ignore_index=-100).to(device)
# criterion=nn.BCELoss(weight=weight,reduction='elementwise_mean').cuda()
check_mkdir(train_args['ckpt_path'])
check_mkdir(os.path.join(train_args['ckpt_path'], args['exp']))
check_mkdir(
os.path.join(train_args['ckpt_path'], args['exp'],
train_args['exp_name']))
open(
os.path.join(train_args['ckpt_path'], args['exp'],
train_args['exp_name'],
str(time.time()) + '.txt'),
'w').write(str(train_args) + '\n\n')
"""###############-------start training--------###############"""
for epoch in range(curr_epoch, train_args['epoch_num'] + 1):
adjust_lr(optimizer, epoch)
train(train_loader, model, criterion, optimizer, epoch, train_args,
device)
val_loss = validate(val_loader, model, criterion, optimizer, restore,
epoch, train_args, visualize, device)
writer.close()
def validate(val_loader, model, criterion, optimizer, restore, epoch,
train_args, visualize, device):
model.eval()
val_loss = AverageMeter()
inputs_all, gts_all, predictions_all = [], [], []
for vi, data in enumerate(val_loader):
inputs, gts, = data
N = inputs.size(0)
inputs = Variable(inputs)
gts = Variable(gts)
if train_args['use_gpu']:
inputs = Variable(inputs).to(device)
gts = Variable(gts).to(device).long()
outputs = model(inputs)
predictions = outputs.data[:].max(1)[1].squeeze_(1).squeeze_(
0).cpu().numpy()
val_loss.update(criterion(outputs, gts).data / N, N)
# val_loss.update(criterion(outputs, gts).data, N)
if random.random() > train_args['val_img_sample_rate']:
inputs_all.append(None)
else:
inputs_all.append(inputs.data[:].squeeze_(0).cpu())
gts_all.append(gts.data[:].squeeze_(0).cpu().numpy())
predictions_all.append(predictions)
acc, acc_cls, mean_iu, fwavacc, single_cls_acc, kappa = evaluate(
predictions_all, gts_all, train_args['num_class'])
if mean_iu > train_args['best_record']['mean_iu']:
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
#
# print(optimizer.param_groups)
snapshot_name = 'epoch_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.7f_kappa_%.5f' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc,
optimizer.param_groups[0]['lr'], kappa)
"""###############-------save models and pred image--------###############"""
if epoch % 30 == 0:
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(
train_args['ckpt_path'], args['exp'],
train_args['exp_name'] + 'pre_img_epoch %d' % epoch)
check_mkdir(to_save_dir)
torch.save(
model.state_dict(),
os.path.join(train_args['ckpt_path'], args['exp'],
train_args['exp_name'], snapshot_name + '.pth'))
torch.save(
optimizer.state_dict(),
os.path.join(train_args['ckpt_path'], args['exp'],
train_args['exp_name'],
'opt_' + snapshot_name + '.pth'))
val_visual = []
for idx, data in enumerate(zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = restore(data[0]) ##反归一化成原图片
gt_pil = dataset_2_cls_4channel.colorize_mask(data[1]) ##上色
predictions_pil = dataset_2_cls_4channel.colorize_mask(data[2])
if train_args['val_save_to_img_file']:
input_pil.save(os.path.join(to_save_dir, '%d_input.png' % idx))
predictions_pil.save(
os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([
visualize(input_pil.convert('RGB')),
visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))
])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=3, padding=5)
writer.add_image(snapshot_name, val_visual)
print(
'--------------------------------starting validating-----------------------------------'
)
print('val_single_cls_acc:')
print(single_cls_acc)
print(
'[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f],[kappa %.5f]'
% (epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, kappa))
print(
'best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d]'
% (train_args['best_record']['val_loss'],
train_args['best_record']['acc'],
train_args['best_record']['acc_cls'],
train_args['best_record']['mean_iu'],
train_args['best_record']['fwavacc'],
train_args['best_record']['epoch']))
print(
'--------------------------------finish validating-------------------------------------'
)
writer.add_scalars(
'single_cls_acc', {
'background': single_cls_acc[0],
train_args['training_cls']: single_cls_acc[1]
}, epoch)
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
writer.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch)
writer.add_scalar('kappa', kappa, epoch)
model.train()
return val_loss.avg
def main():
args = get_args()
check_mkdir('./logs')
logging_name = './logs/{}_{}_lr_{}.txt'.format(args.networks,
args.optimizer, args.lr)
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter('[%(asctime)10s][%(levelname)s] %(message)s',
datefmt='%Y/%m/%d %H:%M:%S')
stream_handler = logging.StreamHandler()
stream_handler.setFormatter(formatter)
file_handler = logging.FileHandler(logging_name)
file_handler.setFormatter(formatter)
logger.addHandler(stream_handler)
logger.addHandler(file_handler)
logger.info('arguments:{}'.format(" ".join(sys.argv)))
if args.ignite is False:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = mobile_hair.MobileMattingFCN()
if torch.cuda.is_available():
if torch.cuda.device_count() > 1:
print('multi gpu')
model = torch.nn.DataParallel(model)
model.to(device)
loss = mobile_hair.HairMattingLoss()
optimizer = get_optimizer(args.optimizer, model, args.lr,
args.momentum)
# torch.optim.Adam(filter(lambda p: p.requires_grad,model.parameters()), lr=0.0001, betas=(0.9, 0.999))
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, 'min')
train_without_ignite(model,
loss,
batch_size=args.batch_size,
img_size=args.img_size,
epochs=args.epochs,
lr=args.lr,
num_workers=args.num_workers,
optimizer=optimizer,
logger=logger,
gray_image=True,
scheduler=scheduler,
viz=args.visdom)
else:
train_with_ignite(networks=args.networks,
dataset=args.dataset,
data_dir=args.data_dir,
batch_size=args.batch_size,
epochs=args.epochs,
lr=args.lr,
num_workers=args.num_workers,
optimizer=args.optimizer,
momentum=args.momentum,
img_size=args.img_size,
logger=logger)
def main():
args = parse_arguments()
# Setting paths as per psc
if args.aws_or_psc == 'psc':
storage_path = 'pylon5/' + args.psc_groupid + '/' + args.psc_userid + '/active-vision-storage/'
elif args.aws_or_psc == 'aws':
storage_path = '../../../active-vision-storage/'
logger = Logger(storage_path + 'av_tblogs', name=args.tb_logdir)
ckpt_path = storage_path + 'ckpt'
exp_name = args.net
image_folder = storage_path + 'av_data/images'
mask_folder = storage_path + 'av_data/masks'
target_folder = storage_path + 'av_data/targets'
if args.net == 'vgg':
net = FCN8s(num_classes=2)
elif args.net == 'alexnet':
net = FCN8s_alex(num_classes=2)
if use_cuda:
net.cuda()
if len(args.snapshot) == 0:
curr_epoch = 1
train_args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
print('training resumes from ' + args.snapshot)
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name, args.snapshot)))
split_snapshot = args.snapshot.split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]),
'fwavacc': float(split_snapshot[11])
}
transform = transforms.Compose([transforms.ToTensor()])
train_image_names, train_mask_names, train_target_names, val_image_names, val_mask_names, val_target_names = make_dataset(
'train', image_folder, mask_folder, target_folder)
train_dataset = active_vision_dataloader.AV(train_image_names,
train_mask_names,
train_target_names,
image_folder,
mask_folder,
target_folder,
mode='train',
transform=transform)
val_dataset = active_vision_dataloader.AV(val_image_names,
val_mask_names,
val_target_names,
image_folder,
mask_folder,
target_folder,
mode='val',
transform=transform)
# Setting Dataloaders
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4,
drop_last=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
num_workers=4,
batch_size=args.batch_size,
drop_last=True,
shuffle=True)
criterion = CrossEntropyLoss2d(size_average=False)
if use_cuda:
criterion.cuda()
optimizer = optim.Adam([{
'params': [
param
for name, param in net.named_parameters() if name[-4:] == 'bias'
],
'lr':
2 * args.lr
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] != 'bias'
],
'lr':
args.lr,
'weight_decay':
args.weight_decay
}],
betas=(args.momentum, 0.999))
if len(args.snapshot) > 0:
optimizer.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, 'opt_' + args.snapshot)))
optimizer.param_groups[0]['lr'] = 2 * args.lr
optimizer.param_groups[1]['lr'] = args.lr
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
# open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(train_args) + '\n\n')
scheduler = ReduceLROnPlateau(optimizer,
'min',
patience=args.lr_patience,
min_lr=1e-10,
verbose=True)
for epoch in range(curr_epoch, args.epoch_num + 1):
train(train_loader, net, criterion, optimizer, epoch, logger)
if epoch % args.valid_freq == 0:
val_loss = validate(val_loader, net, criterion, optimizer, epoch,
logger, ckpt_path, exp_name)
scheduler.step(val_loss)
def validate(val_loader, net, criterion, optimizer, epoch, train_args, restore, visualize):
net.eval()
val_loss = AverageMeter()
inputs_all, gts_all, predictions_all = [], [], []
for vi, data in enumerate(val_loader):
inputs, gts = data
N = inputs.size(0)
inputs = Variable(inputs, volatile=True).cuda()
gts = Variable(gts, volatile=True).cuda()
outputs = net(inputs)
predictions = outputs.data.max(1)[1].squeeze_(1).squeeze_(0).cpu().numpy()
val_loss.update(criterion(outputs, gts).data[0] / N, N)
if random.random() > train_args['val_img_sample_rate']:
inputs_all.append(None)
else:
inputs_all.append(inputs.data.squeeze_(0).cpu())
gts_all.append(gts.data.squeeze_(0).cpu().numpy())
predictions_all.append(predictions)
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all, voc.num_classes)
if mean_iu > train_args['best_record']['mean_iu']:
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
snapshot_name = 'epoch_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.10f' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, optimizer.param_groups[1]['lr']
)
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, str(epoch))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = restore(data[0])
gt_pil = voc.colorize_mask(data[1])
predictions_pil = voc.colorize_mask(data[2])
if train_args['val_save_to_img_file']:
input_pil.save(os.path.join(to_save_dir, '%d_input.png' % idx))
predictions_pil.save(os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([visualize(input_pil.convert('RGB')), visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=3, padding=5)
writer.add_image(snapshot_name, val_visual)
print('--------------------------------------------------------------------')
print('[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc))
print('best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d]' % (
train_args['best_record']['val_loss'], train_args['best_record']['acc'], train_args['best_record']['acc_cls'],
train_args['best_record']['mean_iu'], train_args['best_record']['fwavacc'], train_args['best_record']['epoch']))
print('--------------------------------------------------------------------')
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
writer.add_scalar('lr', optimizer.param_groups[1]['lr'], epoch)
net.train()
return val_loss.avg
def main():
# args = parse_args()
torch.backends.cudnn.benchmark = True
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
# # if args.seed:
# random.seed(args.seed)
# np.random.seed(args.seed)
# torch.manual_seed(args.seed)
# # if args.gpu:
# torch.cuda.manual_seed_all(args.seed)
seed = 63
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# if args.gpu:
torch.cuda.manual_seed_all(seed)
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
# train_transforms = transforms.Compose([
# transforms.RandomCrop(args['crop_size']),
# transforms.RandomRotation(90),
# transforms.RandomHorizontalFlip(p=0.5),
# transforms.RandomVerticalFlip(p=0.5),
# ])
short_size = int(min(args['input_size']) / 0.875)
# val_transforms = transforms.Compose([
# transforms.Scale(short_size, interpolation=Image.NEAREST),
# # joint_transforms.Scale(short_size),
# transforms.CenterCrop(args['input_size'])
# ])
train_joint_transform = joint_transforms.Compose([
# joint_transforms.Scale(short_size),
joint_transforms.RandomCrop(args['crop_size']),
joint_transforms.RandomHorizontallyFlip(),
joint_transforms.RandomRotate(90)
])
val_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(short_size),
joint_transforms.CenterCrop(args['input_size'])
])
input_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(*mean_std)])
target_transform = extended_transforms.MaskToTensor()
restore_transform = transforms.Compose(
[extended_transforms.DeNormalize(*mean_std),
transforms.ToPILImage()])
visualize = transforms.ToTensor()
train_set = cityscapes.CityScapes('train',
joint_transform=train_joint_transform,
transform=input_transform,
target_transform=target_transform)
# train_set = cityscapes.CityScapes('train', transform=train_transforms)
train_loader = DataLoader(train_set,
batch_size=args['train_batch_size'],
num_workers=8,
shuffle=True)
val_set = cityscapes.CityScapes('val',
joint_transform=val_joint_transform,
transform=input_transform,
target_transform=target_transform)
# val_set = cityscapes.CityScapes('val', transform=val_transforms)
val_loader = DataLoader(val_set,
batch_size=args['val_batch_size'],
num_workers=8,
shuffle=True)
print(len(train_loader), len(val_loader))
# sdf
vgg_model = VGGNet(requires_grad=True, remove_fc=True)
net = FCN8s(pretrained_net=vgg_model,
n_class=cityscapes.num_classes,
dropout_rate=0.4)
# net.apply(init_weights)
criterion = nn.CrossEntropyLoss(ignore_index=cityscapes.ignore_label)
optimizer = optim.Adam(net.parameters(), lr=1e-4)
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(
os.path.join(ckpt_path, exp_name,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(args) + '\n\n')
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, 'min', patience=args['lr_patience'], min_lr=1e-10)
vgg_model = vgg_model.to(device)
net = net.to(device)
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
if len(args['snapshot']) == 0:
curr_epoch = 1
args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0
}
else:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9][:-4])
}
criterion.to(device)
for epoch in range(curr_epoch, args['epoch_num'] + 1):
train(train_loader, net, device, criterion, optimizer, epoch, args)
val_loss = validate(val_loader, net, device, criterion, optimizer,
epoch, args, restore_transform, visualize)
scheduler.step(val_loss)
lr = args.lr * 0.01
if epoch >= 100:
lr = args.lr * 0.001
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if __name__ == "__main__":
args = parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
gpu_num = max(len(args.gpu_id.split(',')), 1)
model_name = 'resnet18'
log_dir = "logs/%s_%s" % (time.strftime("%b%d-%H%M",
time.localtime()), model_name)
check_mkdir(log_dir)
log = Logger(log_dir + '/train.log')
log.print(args)
device = torch.device('cuda')
model = ResNet18().to(device)
model = nn.DataParallel(model, device_ids=[i for i in range(gpu_num)])
train_loader, test_loader = prepare_cifar(args.batch_size,
args.test_batch_size)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
best_epoch, best_robust_acc = 0, 0.
def main(model_name):
# TODO: parse args.
n_classes = 19
#batch_size = 2
batch_size = 1 #24
n_workers = 12
n_semantic_pretrain = 0 # 500 # First train only on semantics.
n_epochs = 500
validation_step = 15
# TODO: implement resize as pil_transform
resize = None # (256, 512)
cityscapes_directory = "/home/<someuser>/cityscapes"
output_directory = "tmp/"
# Uncomment next line when you've set all directories.
raise ValueError("Please set the input/output directories.")
checkpoint = None
#checkpoint = (
# "weights/...pth",
# <fill in epoch>)
# --- Setup loss functions.
classification_loss = nn.CrossEntropyLoss(ignore_index=255)
regression_loss = nn.MSELoss(reduction='elementwise_mean')
print("--- Load model.")
if model_name == 'DRNRegressionDownsampled':
classification_loss = None
regression_loss = nn.MSELoss(reduction='elementwise_mean')
dataset_kwargs = {
'pil_transforms':
None,
'gt_pil_transforms': [ModeDownsample(8)],
'fit_gt_pil_transforms':
[transforms.Resize(size=(784 // 8, 1792 // 8), interpolation=2)],
'input_transforms': [
transforms.Normalize(mean=[0.290101, 0.328081, 0.286964],
std=[0.182954, 0.186566, 0.184475])
],
'tensor_transforms':
None
}
model = DRNRegressionDownsampled(
model_name='drn_d_22',
classes=n_classes,
pretrained_dict=torch.load('./weights/drn_d_22_cityscapes.pth'))
model.cuda()
parameters = model.parameters()
else:
raise ValueError("Model \"{}\" not found!".format(model_name))
optimizer = optim.Adam(parameters)
start_epoch = 0
if checkpoint is not None:
print("Loading from checkpoint {}".format(checkpoint))
model.load_state_dict(torch.load(checkpoint[0]))
optimizer.load_state_dict(torch.load(checkpoint[1]))
start_epoch = checkpoint[2] + 1
print("--- Setup dataset and dataloaders.")
train_set = Cityscapes(data_split='subtrain',
cityscapes_directory=cityscapes_directory,
**dataset_kwargs)
train_loader = data.DataLoader(train_set,
batch_size=batch_size,
num_workers=n_workers,
shuffle=True)
val_set = Cityscapes(data_split='subtrainval',
cityscapes_directory=cityscapes_directory,
**dataset_kwargs)
val_loader = data.DataLoader(val_set,
batch_size=batch_size,
num_workers=n_workers)
# Sample 10 validation indices for visualization.
#validation_idxs = np.random.choice(np.arange(len(val_set)),
# size=min(9, len(val_set)),
# replace=False)
# Nah, let's pick them ourselves for now.
validation_idxs = [17, 241, 287, 304, 123, 458, 1, 14, 139, 388]
if True:
print("--- Setup visual validation.")
# Save them for comparison.
check_mkdir('{}/validationimgs'.format(output_directory))
check_mkdir('{}/offsets_gt'.format(output_directory))
check_mkdir('{}/semantic_gt'.format(output_directory))
for validation_idx in validation_idxs:
img_pil, _, _ = val_set.load_fit_gt_PIL_images(validation_idx)
img, semantic_gt, offset_gt = val_set[validation_idx]
img_pil.save("{}/validationimgs/id{:03}.png".format(
output_directory, validation_idx))
visualize_semantics(
img_pil, semantic_gt,
"{}/semantic_gt/id{:03}".format(output_directory,
validation_idx))
visualize_positionplusoffset(offset_gt,
"{}/offsets_gt/id{:03}_mean".format(
output_directory, validation_idx),
groundtruth=offset_gt)
visualize_offsethsv(
offset_gt,
"{}/offsets_gt/id{:03}".format(output_directory,
validation_idx))
print("--- Training.")
rlosses = []
closses = []
for epoch in range(start_epoch, n_epochs):
model.train()
total_rloss = 0
total_closs = 0
for batch_idx, batch_data in enumerate(train_loader):
img = batch_data[0].cuda()
semantic_gt = batch_data[1].cuda()
instance_offset_gt = batch_data[2].cuda()
del batch_data
optimizer.zero_grad()
outputs = model(img)
batch_rloss = 0
batch_closs = 0
loss = 0
closs = 0
rloss = 0
if regression_loss is not None:
predicted_offset = outputs[:, -2:]
rloss = regression_loss(predicted_offset, instance_offset_gt)
batch_rloss += int(rloss.detach().cpu())
total_rloss += batch_rloss
loss += rloss
if classification_loss is not None:
closs = classification_loss(outputs[:, :n_classes],
semantic_gt)
batch_closs += int(closs.detach().cpu())
total_closs += batch_closs
loss += closs
loss.backward()
optimizer.step()
if batch_idx % 30 == 0 and batch_idx != 0:
print('\t[batch {}/{}], [batch mean - closs {:5}, rloss {:5}]'.
format(batch_idx, len(train_loader),
batch_closs / img.size(0),
batch_rloss / img.size(0)))
del img, semantic_gt, instance_offset_gt, outputs, rloss, closs, loss
total_closs /= len(train_set)
total_rloss /= len(train_set)
print('[epoch {}], [mean train - closs {:5}, rloss {:5}]'.format(
epoch, total_closs, total_rloss))
rlosses.append(total_rloss)
closses.append(total_closs)
plt.plot(np.arange(start_epoch, epoch + 1), rlosses)
plt.savefig('{}/rlosses.svg'.format(output_directory))
plt.close('all')
plt.plot(np.arange(start_epoch, epoch + 1), closses)
plt.savefig('{}/closses.svg'.format(output_directory))
plt.close('all')
plt.plot(np.arange(start_epoch, epoch + 1), np.add(rlosses, closses))
plt.savefig('{}/losses.svg'.format(output_directory))
plt.close('all')
# --- Visual validation.
if (epoch % validation_step) == 0:
# Save model parameters.
check_mkdir('{}/models'.format(output_directory))
torch.save(
model.state_dict(),
'{}/models/Net_epoch{}.pth'.format(output_directory, epoch))
torch.save(
optimizer.state_dict(),
'{}/models/Adam_epoch{}.pth'.format(output_directory, epoch))
# Visualize validation imgs.
check_mkdir('{}/offsets'.format(output_directory))
check_mkdir('{}/offsets/means'.format(output_directory))
check_mkdir('{}/semantics'.format(output_directory))
check_mkdir('{}/semantics/overlay'.format(output_directory))
model.eval()
for validation_idx in validation_idxs:
img_pil, _, _ = val_set.load_PIL_images(validation_idx)
img, _, offset_gt = val_set[validation_idx]
img = img.unsqueeze(0).cuda()
with torch.no_grad():
outputs = model(img)
epoch_filename = 'id{:03}_epoch{:05}'\
.format(validation_idx, epoch)
if classification_loss is not None:
visualize_semantics(
img_pil, outputs,
"{}/semantics/{}".format(output_directory,
epoch_filename),
"{}/semantics/overlay/{}".format(
output_directory, epoch_filename))
if regression_loss is not None:
visualize_offsethsv(
outputs.detach(),
"{}/offsets/{}".format(output_directory,
epoch_filename))
visualize_positionplusoffset(outputs,
"{}/offsets/means/{}".format(
output_directory,
epoch_filename),
groundtruth=offset_gt)
def test(test_loader, net, criterion, epoch, num_known_classes, num_unknown_classes, hidden, train_args, save_images, save_model, weibull_model):
# Setting network for evaluation mode.
net.eval()
with torch.no_grad():
# Creating output directory.
check_mkdir(os.path.join(outp_path, exp_name, 'epoch_' + str(epoch)))
# Iterating over batches.
for i, data in enumerate(test_loader):
print('Test Batch %d/%d' % (i + 1, len(test_loader)))
# Obtaining images, labels and paths for batch.
inps_batch, labs_batch, true_batch, img_name = data
inps_batch = inps_batch.squeeze()
labs_batch = labs_batch.squeeze()
true_batch = true_batch.squeeze()
# Iterating over patches inside batch.
for j in range(inps_batch.size(0)):
print(' Test MiniBatch %d/%d' % (j + 1, inps_batch.size(0)))
tic = time.time()
for k in range(inps_batch.size(1)):
print(' Test MiniMiniBatch %d/%d' % (k + 1, inps_batch.size(1)))
sys.stdout.flush()
inps = inps_batch[j, k].unsqueeze(0)
labs = labs_batch[j, k].unsqueeze(0)
true = true_batch[j, k].unsqueeze(0)
# Casting tensors to cuda.
inps, labs, true = inps.cuda(args['device']), labs.cuda(args['device']), true.cuda(args['device'])
# Casting to cuda variables.
inps = Variable(inps).cuda(args['device'])
labs = Variable(labs).cuda(args['device'])
true = Variable(true).cuda(args['device'])
# Forwarding.
outs = net(inps)
# Computing loss.
soft_outs = F.softmax(outs, dim=1)
open_outs = recalibrate_scores(weibull_model,
outs.permute(0, 2, 3, 1).cpu().numpy(),
soft_outs.permute(0, 2, 3, 1).cpu().numpy(),
num_known_classes,
alpharank=num_known_classes,
distance_type=args['distance_type'])
open_outs = np.asarray(open_outs)
open_outs = open_outs.reshape(outs.size(0), outs.size(2), outs.size(3), open_outs.shape[1])
# Obtaining predictions.
prds = open_outs.argmax(axis=3)
prds[open_outs.max(axis=3) < args['open_threshold']] = num_known_classes
# Saving predictions.
if (save_images):
pred_path = os.path.join(outp_path, exp_name, 'epoch_' + str(epoch), img_name[0].replace('.tif', '_pred_' + str(j) + '_' + str(k) + '.png'))
prob_path = os.path.join(outp_path, exp_name, 'epoch_' + str(epoch), img_name[0].replace('.tif', '_prob_' + str(j) + '_' + str(k) + '.npy'))
io.imsave(pred_path, util.img_as_ubyte(prds.squeeze()))
np.save(prob_path, np.transpose(open_outs.squeeze(), (2, 0, 1)))
toc = time.time()
print(' Elapsed Time: %.2f' % (toc - tic))
sys.stdout.flush()
def main():
torch.backends.cudnn.benchmark = True
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
vgg_model = VGGNet(requires_grad=True, remove_fc=True)
net = FCN8s(pretrained_net=vgg_model,
n_class=cityscapes.num_classes,
dropout_rate=0.4)
print('load model ' + args['snapshot'])
vgg_model = vgg_model.to(device)
net = net.to(device)
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
net.load_state_dict(
torch.load(os.path.join(ckpt_path, args['exp_name'],
args['snapshot'])))
net.eval()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
short_size = int(min(args['input_size']) / 0.875)
val_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(short_size),
joint_transforms.CenterCrop(args['input_size'])
])
test_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(*mean_std)])
target_transform = extended_transforms.MaskToTensor()
restore_transform = transforms.Compose(
[extended_transforms.DeNormalize(*mean_std),
transforms.ToPILImage()])
# test_set = cityscapes.CityScapes('test', transform=test_transform)
test_set = cityscapes.CityScapes('test',
joint_transform=val_joint_transform,
transform=test_transform,
target_transform=target_transform)
test_loader = DataLoader(test_set,
batch_size=1,
num_workers=8,
shuffle=False)
transform = transforms.ToPILImage()
check_mkdir(os.path.join(ckpt_path, args['exp_name'], 'test'))
gts_all, predictions_all = [], []
count = 0
for vi, data in enumerate(test_loader):
# img_name, img = data
img_name, img, gts = data
img_name = img_name[0]
# print(img_name)
img_name = img_name.split('/')[-1]
# img.save(os.path.join(ckpt_path, args['exp_name'], 'test', img_name))
img_transform = restore_transform(img[0])
# img_transform = img_transform.convert('RGB')
img_transform.save(
os.path.join(ckpt_path, args['exp_name'], 'test', img_name))
img_name = img_name.split('_leftImg8bit.png')[0]
# img = Variable(img, volatile=True).cuda()
img, gts = img.to(device), gts.to(device)
output = net(img)
prediction = output.data.max(1)[1].squeeze_(1).squeeze_(
0).cpu().numpy()
prediction_img = cityscapes.colorize_mask(prediction)
# print(type(prediction_img))
prediction_img.save(
os.path.join(ckpt_path, args['exp_name'], 'test',
img_name + '.png'))
# print(ckpt_path, args['exp_name'], 'test', img_name + '.png')
print('%d / %d' % (vi + 1, len(test_loader)))
gts_all.append(gts.data.cpu().numpy())
predictions_all.append(prediction)
# break
# if count == 1:
# break
# count += 1
gts_all = np.concatenate(gts_all)
predictions_all = np.concatenate(prediction)
acc, acc_cls, mean_iou, _ = evaluate(predictions_all, gts_all,
cityscapes.num_classes)
print(
'-----------------------------------------------------------------------------------------------------------'
)
print('[acc %.5f], [acc_cls %.5f], [mean_iu %.5f]' %
(acc, acc_cls, mean_iu))
momentum = 0
log_interval = 12
train_size = 100
stat = defaultdict(dict)
stat['n_epochs'] = 40
stat['bsize'] = 2
stat['iterations'] = 8 #num for itrs for couplings updates
stat['weight_decay'] = 5e-4
stat['dev'] = 'cuda: 2' if torch.cuda.is_available() else 'cpu'
# In[4]:
#saving path
MNIST_tran_ini = './CIFAR_initialstatus'
stat['savingmodel'] = './CIFAR_stat'
check_mkdir(stat['savingmodel'])
check_mkdir(MNIST_tran_ini)
# In[5]:
##animal and vehicle dataset
for i in range(2):
if i == 0:
index = [2, 3, 4, 5, 6, 7]
else:
index = [0, 1, 8, 9]
train = thv.datasets.CIFAR10(root='./data',
train=True,
download=True,
transform=transform_train)
train.targets = torch.tensor(train.targets)
inputs = Variable(inputs.cuda(),volatile=True)
outputs = model(inputs)
outputs = F.upsample(outputs, scale_factor=2, mode='bilinear')
pred = outputs.data.max(1)[1]
for predid in range(18, -1, -1):
pred[pred==predid] = class_map[predid]
prediction = pred.cpu().squeeze_().float().numpy()
prediction = Image.fromarray(prediction).convert("L")
prediction.save(os.path.join(test_args.save_path, filepath[0]))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Hyperparams')
parser.add_argument('--load_param', nargs='?', type=str, default=None,
help='Path to pretrained parameters to test on')
parser.add_argument('--save_path', nargs='?', type=str, default=None,
help='Path to save predicted images')
args = parser.parse_args()
check_mkdir(args.save_path)
main(args)
def main():
net = FCN8s(num_classes=cityscapes.num_classes).cuda()
if len(args['snapshot']) == 0:
curr_epoch = 1
args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
print 'training resumes from ' + args['snapshot']
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]),
'fwavacc': float(split_snapshot[11])
}
net.train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
short_size = int(min(args['input_size']) / 0.875)
train_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(short_size),
joint_transforms.RandomCrop(args['input_size']),
joint_transforms.RandomHorizontallyFlip()
])
val_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(short_size),
joint_transforms.CenterCrop(args['input_size'])
])
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage()
])
visualize = standard_transforms.ToTensor()
train_set = cityscapes.CityScapes('fine',
'train',
joint_transform=train_joint_transform,
transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=args['train_batch_size'],
num_workers=8,
shuffle=True)
val_set = cityscapes.CityScapes('fine',
'val',
joint_transform=val_joint_transform,
transform=input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=args['val_batch_size'],
num_workers=8,
shuffle=False)
criterion = CrossEntropyLoss2d(
size_average=False, ignore_index=cityscapes.ignore_label).cuda()
optimizer = optim.SGD([{
'params': [
param
for name, param in net.named_parameters() if name[-4:] == 'bias'
],
'lr':
2 * args['lr']
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] != 'bias'
],
'lr':
args['lr'],
'weight_decay':
args['weight_decay']
}],
momentum=args['momentum'])
if len(args['snapshot']) > 0:
optimizer.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, 'opt_' + args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * args['lr']
optimizer.param_groups[1]['lr'] = args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(
os.path.join(ckpt_path, exp_name,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(args) + '\n\n')
scheduler = ReduceLROnPlateau(optimizer,
'min',
patience=args['lr_patience'],
min_lr=1e-10)
for epoch in range(curr_epoch, args['epoch_num'] + 1):
train(train_loader, net, criterion, optimizer, epoch, args)
val_loss = validate(val_loader, net, criterion, optimizer, epoch, args,
restore_transform, visualize)
scheduler.step(val_loss)
def main(train_args):
net = PSPNet(num_classes=voc.num_classes).cuda()
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0}
else:
print ('training resumes from ' + train_args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])}
net.train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_simul_transform = simul_transforms.Compose([
simul_transforms.RandomSized(train_args['input_size']),
simul_transforms.RandomRotate(10),
simul_transforms.RandomHorizontallyFlip()
])
val_simul_transform = simul_transforms.Scale(train_args['input_size'])
train_input_transform = standard_transforms.Compose([
extended_transforms.RandomGaussianBlur(),
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage(),
])
visualize = standard_transforms.Compose([
standard_transforms.Scale(400),
standard_transforms.CenterCrop(400),
standard_transforms.ToTensor()
])
train_set = voc.VOC('train', joint_transform=train_simul_transform, transform=train_input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=train_args['train_batch_size'], num_workers=8, shuffle=True)
val_set = voc.VOC('val', joint_transform=val_simul_transform, transform=val_input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set, batch_size=1, num_workers=8, shuffle=False)
criterion = CrossEntropyLoss2d(size_average=True, ignore_index=voc.ignore_label).cuda()
optimizer = optim.SGD([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * train_args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': train_args['lr'], 'weight_decay': train_args['weight_decay']}
], momentum=train_args['momentum'])
if len(train_args['snapshot']) > 0:
optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + train_args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * train_args['lr']
optimizer.param_groups[1]['lr'] = train_args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(train_args) + '\n\n')
train(train_loader, net, criterion, optimizer, curr_epoch, train_args, val_loader, restore_transform, visualize)
def main(train_args):
import pdb
pdb.set_trace()
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
net = FCN8s(num_classes=plant.num_classes).cuda()
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
print('training resumes from ' + train_args['snapshot'])
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]),
'fwavacc': float(split_snapshot[11])
}
net.train()
mean_std = ([0.385, 0.431, 0.452], [0.289, 0.294, 0.285])
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage(),
])
visualize = standard_transforms.Compose([
standard_transforms.Scale(500),
standard_transforms.CenterCrop(500),
standard_transforms.ToTensor()
])
data_aug = Compose([RandomRotate(10), RandomHorizontallyFlip()])
train_set = plant.Plant('train',
augmentations=data_aug,
transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=4,
num_workers=2,
shuffle=True)
val_set = plant.Plant('val',
transform=input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=1,
num_workers=2,
shuffle=False)
weights = torch.FloatTensor(cfg.train_weights)
criterion = CrossEntropyLoss2d(weight=weights,
size_average=False,
ignore_index=plant.ignore_label).cuda()
optimizer = optim.Adam([{
'params': [
param
for name, param in net.named_parameters() if name[-4:] == 'bias'
],
'lr':
2 * train_args['lr']
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] != 'bias'
],
'lr':
train_args['lr'],
'weight_decay':
train_args['weight_decay']
}],
betas=(train_args['momentum'], 0.999))
if len(train_args['snapshot']) > 0:
optimizer.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name,
'opt_' + train_args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * train_args['lr']
optimizer.param_groups[1]['lr'] = train_args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(
os.path.join(ckpt_path, exp_name,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(train_args) + '\n\n')
#train_args['best_record']['mean_iu'] = 0.50
#curr_epoch = 100
scheduler = ReduceLROnPlateau(optimizer,
'min',
patience=train_args['lr_patience'],
min_lr=1e-10,
verbose=True)
for epoch in range(curr_epoch, train_args['epoch_num'] + 1):
val_loss = validate(val_loader, net, criterion, optimizer, epoch,
train_args, restore_transform, visualize)
train(train_loader, net, criterion, optimizer, epoch, train_args)
scheduler.step(val_loss)
def validate(val_loader, net, criterion, optimizer, epoch, train_args, restore, visualize):
net.eval()
val_loss = AverageMeter()
inputs_all, gts_all, predictions_all = [], [], []
for vi, data in enumerate(val_loader):
inputs, gts = data
N = inputs.size(0)
inputs = Variable(inputs, volatile=True).cuda()
gts = Variable(gts, volatile=True).cuda()
outputs = net(inputs)
predictions = outputs.data.max(1)[1].squeeze_(1).squeeze_(0).cpu().numpy()
val_loss.update(criterion(outputs, gts).data[0] / N, N)
if random.random() > train_args['val_img_sample_rate']:
inputs_all.append(None)
else:
inputs_all.append(inputs.data.squeeze_(0).cpu())
gts_all.append(gts.data.squeeze_(0).cpu().numpy())
predictions_all.append(predictions)
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all, voc.num_classes)
if mean_iu > train_args['best_record']['mean_iu']:
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
snapshot_name = 'epoch_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.10f' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, optimizer.param_groups[1]['lr']
)
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, str(epoch))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = restore(data[0])
gt_pil = voc.colorize_mask(data[1])
predictions_pil = voc.colorize_mask(data[2])
if train_args['val_save_to_img_file']:
input_pil.save(os.path.join(to_save_dir, '%d_input.png' % idx))
predictions_pil.save(os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([visualize(input_pil.convert('RGB')), visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=3, padding=5)
writer.add_image(snapshot_name, val_visual)
print('--------------------------------------------------------------------')
print('[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc))
print('best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d]' % (
train_args['best_record']['val_loss'], train_args['best_record']['acc'], train_args['best_record']['acc_cls'],
train_args['best_record']['mean_iu'], train_args['best_record']['fwavacc'], train_args['best_record']['epoch']))
print('--------------------------------------------------------------------')
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
writer.add_scalar('lr', optimizer.param_groups[1]['lr'], epoch)
net.train()
return val_loss.avg
def validate(epoch, val_loader, sr_model, seg_model, sr_criterion,
psnr_criterion, seg_criterion, sr_optimizer, seg_optimizer):
sr_model.eval()
seg_model.eval()
epoch_loss = 0
epoch_sr_loss = 0
epoch_seg_loss = 0
avg_psnr = 0
avg_psnr_bicubic = 0
msks_all, prds_all = [], []
save_img_epoch = 50
if epoch % save_img_epoch == 0:
check_mkdir(
os.path.join('outputs', exp_name, 'segmentation',
'epoch_' + str(epoch)))
check_mkdir(
os.path.join('outputs', exp_name, 'super-resolution',
'epoch_' + str(epoch)))
for iteration, batch in enumerate(val_loader, 1):
with torch.no_grad():
lr, bicubic, img, msk, img_name = Variable(batch[0]), Variable(
batch[1]), Variable(batch[2]), Variable(batch[3]), batch[4]
if cuda:
lr = lr.cuda(gpus_list[0])
img = img.cuda(gpus_list[0])
bicubic = bicubic.cuda(gpus_list[0])
msk = msk.cuda(gpus_list[0])
sr_optimizer.zero_grad()
seg_optimizer.zero_grad()
with torch.no_grad():
sr_prediction = sr_model(lr)
if opt.sr_residual:
sr_prediction = sr_prediction + bicubic
sr_loss = sr_criterion(sr_prediction, img)
mse = psnr_criterion(sr_prediction, img)
psnr = 10 * log10(1 / mse.data)
avg_psnr += psnr
mse = psnr_criterion(bicubic, img)
psnr = 10 * log10(1 / mse.data)
avg_psnr_bicubic += psnr
seg_prediction = seg_model(sr_prediction)
#seg_prediction = seg_model(sr_prediction*255)
msk = msk.long()
seg_loss = seg_criterion(seg_prediction, msk)
seg_prediction = seg_prediction.data.max(1)[1].squeeze_(
1).squeeze_(0).cpu().numpy()
msks_all.append(msk.squeeze_(0).cpu().numpy())
prds_all.append(seg_prediction)
total_loss = opt.alfa * sr_loss + opt.beta * seg_loss
epoch_sr_loss += sr_loss.data
epoch_seg_loss += seg_loss.data
epoch_loss += total_loss.data
if epoch % save_img_epoch == 0:
tmp_path_seg = os.path.join('outputs', exp_name,
'segmentation',
'epoch_{}'.format(epoch),
img_name[0])
tmp_path_sr = os.path.join('outputs', exp_name,
'super-resolution',
'epoch_{}'.format(epoch),
img_name[0])
if 'grss' in opt.data_dir:
save_img_seg(seg_prediction, tmp_path_seg, msk)
else:
save_img_seg(seg_prediction, tmp_path_seg)
save_img_sr(sr_prediction, tmp_path_sr)
acc, acc_cls, mean_iou, iou, fwavacc, kappa = evaluate(
prds_all, msks_all, opt.num_classes)
#print('--------------------------------------------------------------------')
print(
'VAL: [epoch %d], [loss %.4f], [sr_loss %.4f], [seg_loss %.4f], [PSNR %.4f], [acc %.4f], [acc_cls %.4f], [iou %.4f], [fwavacc %.4f], [kappa %.4f]'
% (epoch, epoch_loss / len(val_loader),
epoch_sr_loss / len(val_loader), epoch_seg_loss / len(val_loader),
avg_psnr / len(val_loader), acc, acc_cls, mean_iou, fwavacc, kappa))
print('Bicubic PSNR: %.4f' % (avg_psnr_bicubic / len(val_loader)))
print(
'--------------------------------------------------------------------')
return mean_iou
def main(train_args):
net = FCN8s(num_classes=voc.num_classes, caffe=True).cuda()
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0}
else:
print('training resumes from ' + train_args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])}
net.train()
mean_std = ([103.939, 116.779, 123.68], [1.0, 1.0, 1.0])
input_transform = standard_transforms.Compose([
extended_transforms.FlipChannels(),
standard_transforms.ToTensor(),
standard_transforms.Lambda(lambda x: x.mul_(255)),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.Lambda(lambda x: x.div_(255)),
standard_transforms.ToPILImage(),
extended_transforms.FlipChannels()
])
visualize = standard_transforms.Compose([
standard_transforms.Scale(400),
standard_transforms.CenterCrop(400),
standard_transforms.ToTensor()
])
train_set = voc.VOC('train', transform=input_transform, target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=1, num_workers=4, shuffle=True)
val_set = voc.VOC('val', transform=input_transform, target_transform=target_transform)
val_loader = DataLoader(val_set, batch_size=1, num_workers=4, shuffle=False)
criterion = CrossEntropyLoss2d(size_average=False, ignore_index=voc.ignore_label).cuda()
optimizer = optim.SGD([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * train_args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': train_args['lr'], 'weight_decay': train_args['weight_decay']}
], momentum=train_args['momentum'])
if len(train_args['snapshot']) > 0:
optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + train_args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * train_args['lr']
optimizer.param_groups[1]['lr'] = train_args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(train_args) + '\n\n')
scheduler = ReduceLROnPlateau(optimizer, 'min', patience=train_args['lr_patience'], min_lr=1e-10, verbose=True)
for epoch in range(curr_epoch, train_args['epoch_num'] + 1):
train(train_loader, net, criterion, optimizer, epoch, train_args)
val_loss = validate(val_loader, net, criterion, optimizer, epoch, train_args, restore_transform, visualize)
scheduler.step(val_loss)
def validate(val_loader, net, criterion, optimizer, epoch, iter_num, train_args, visualize, val_set=None):
# the following code is written assuming that batch size is 1
net.eval()
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, '%d_%d' % (epoch, iter_num))
check_mkdir(to_save_dir)
outs=[]
val_loss = AverageMeter()
gts_all = np.zeros((len(val_loader), int(args['short_size']), int(args['longer_size'])), dtype=int)
predictions_all = np.zeros((len(val_loader), int(args['short_size']), int(args['longer_size'])), dtype=int)
#gts_all = np.zeros((len(val_loader), int(args['longer_size']), int(args['short_size'])), dtype=int)
#predictions_all = np.zeros((len(val_loader), int(args['longer_size']), int(args['short_size'])), dtype=int)
for vi, data in enumerate(tqdm(val_loader)):
input, gt, slices_info = data
assert len(input.size()) == 5 and len(gt.size()) == 4 and len(slices_info.size()) == 3
input.transpose_(0, 1)
gt.transpose_(0, 1)
slices_info.squeeze_(0)
assert input.size()[3:] == gt.size()[2:]
#count = torch.zeros(int(args['longer_size']), int(args['short_size'])).cuda()
#output = torch.zeros(num_classes, int(args['longer_size']), int(args['short_size'])).cuda()
count = torch.zeros(int(args['short_size']), int(args['longer_size'])).cuda()
output = torch.zeros(num_classes, int(args['short_size']), int(args['longer_size'])).cuda()
slice_batch_pixel_size = input.size(1) * input.size(3) * input.size(4)
for input_slice, gt_slice, info in zip(input, gt, slices_info):
# print(gt_slice.cpu().numpy())
# print(np.amax(gt_slice.cpu().numpy()))
gt_slice = Variable(gt_slice, volatile=True).cuda()
input_slice = Variable(input_slice, volatile=True).cuda()
output_slice = net(input_slice)
assert output_slice.size()[2:] == gt_slice.size()[1:]
assert output_slice.size()[1] == num_classes
# print(output_slice.size())
# print(output.size())
output[:, info[0]: info[1], info[2]: info[3]] += output_slice[0, :, :info[4], :info[5]].data
gts_all[vi, info[0]: info[1], info[2]: info[3]] += gt_slice[0, :info[4], :info[5]].data.cpu().numpy()
count[info[0]: info[1], info[2]: info[3]] += 1
val_loss.update(criterion(output_slice, gt_slice).data[0], slice_batch_pixel_size)
output /= count
outs.append(output.cpu().numpy()[0,:,:])
np.floor_divide(gts_all[vi, :, :],count.cpu().numpy().astype(int), out=gts_all[vi, :, :])
#gts_all[vi, :, :] /= count.cpu().numpy().astype(int)
plt.figure(1, figsize=(11,11), dpi=100)
plt.subplot(1,3,1)
plt.imshow(output.cpu().numpy()[0,:,:])
plt.subplot(1,3,2)
plt.imshow(output.cpu().numpy()[1, :, :])
plt.subplot(1, 3, 3)
plt.imshow(output.cpu().numpy()[2,:,:])
plt.savefig(os.path.join(to_save_dir, '%d_prediction_imshow.png' % vi), bbox_inches="tight")
plt.close()
img_name = os.path.split(val_set.imgs[vi][0])[-1]
savemat(os.path.join(to_save_dir, img_name[:-4] + ".mat"), {"maps":output.cpu().numpy()})
#temp_out = output.cpu().numpy()
#temp_out[1, :, :] *= 1.5
#print(np.argmax(temp_out, 0))
#predictions_all[vi, :, :] = np.argmax(temp_out, 0)
predictions_all[vi, :, :] = output.max(0)[1].squeeze_(0).cpu().numpy()
# print('validating: %d / %d' % (vi + 1, len(val_loader)))
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all, num_classes)
if val_loss.avg < train_args['best_record']['val_loss']:
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['iter'] = iter_num
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
snapshot_name = 'epoch_%d_iter_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.10f' % (
epoch, iter_num, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, optimizer.param_groups[1]['lr'])
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
val_visual = []
for idx, data in enumerate(zip(gts_all, predictions_all)):
gt_pil = colorize_mask(data[0])
predictions_pil = colorize_mask(data[1])
if train_args['val_save_to_img_file']:
plt.imshow(outs[idx], cmap="jet")
img_name = os.path.split(val_set.imgs[idx][0])[-1]
img = cv2.imread(val_set.imgs[idx][0], cv2.IMREAD_COLOR)
shutil.copy(val_set.imgs[idx][0], os.path.join(to_save_dir, os.path.split(val_set.imgs[idx][0])[-1]))
predictions_pil.save(os.path.join(to_save_dir, img_name[:-4] + "_prediction.png"))
gt_pil.save(os.path.join(to_save_dir, img_name[:-4] + "_ground_truth.png"))
#cv2.imwrite(os.path.join(to_save_dir, img_name[:-4] + "_gt_vs_pred.png"),
# show_segmentation_into_original_image(img, data[1]))
val_visual.extend([visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))])
val_visual = torch.stack(val_visual, 0)
val_visual = utils.make_grid(val_visual, nrow=2, padding=5)
writer.add_image(snapshot_name, val_visual)
# print('-----------------------------------------------------------------------------------------------------------')
# print('[epoch %d], [iter %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' % (
# epoch, iter_num, val_loss.avg, acc, acc_cls, mean_iu, fwavacc))
#
# print('best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d], '
# '[iter %d]' % (train_args['best_record']['val_loss'], train_args['best_record']['acc'],
# train_args['best_record']['acc_cls'], train_args['best_record']['mean_iu'],
# train_args['best_record']['fwavacc'], train_args['best_record']['epoch'],
# train_args['best_record']['iter']))
#
# print('-----------------------------------------------------------------------------------------------------------')
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
net.train()
return val_loss.avg
def main(train_args):
# weight init
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
torch.nn.init.normal(m.weight.data, mean=0, std=0.01)
torch.nn.init.constant(m.bias.data, 0)
net = VGG(num_classes=VOC.num_classes)
net.apply(weights_init)
net_dict = net.state_dict()
pretrain = torch.load('./vgg16_20M.pkl')
pretrain_dict = pretrain.state_dict()
pretrain_dict = {
'features.' + k: v
for k, v in pretrain_dict.items() if 'features.' + k in net_dict
}
net_dict.update(pretrain_dict)
net.load_state_dict(net_dict)
net = nn.DataParallel(net)
net = net.cuda()
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
print('training resumes from ' + train_args['snapshot'])
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]),
'fwavacc': float(split_snapshot[11])
}
net.train()
mean_std = ([0.408, 0.457, 0.481], [1, 1, 1])
joint_transform_train = joint_transforms.Compose(
[joint_transforms.RandomCrop((321, 321))])
joint_transform_test = joint_transforms.Compose(
[joint_transforms.RandomCrop((512, 512))])
input_transform = standard_transforms.Compose([
#standard_transforms.Resize((321,321)),
#standard_transforms.RandomCrop(224),
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = standard_transforms.Compose([
#standard_transforms.Resize((224,224)),
extended_transforms.MaskToTensor()
])
#target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage(),
])
visualize = standard_transforms.Compose([
standard_transforms.Resize(400),
standard_transforms.CenterCrop(400),
standard_transforms.ToTensor()
])
train_set = VOC.VOC('train',
joint_transform=joint_transform_train,
transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=20,
num_workers=4,
shuffle=True)
val_set = VOC.VOC('val',
joint_transform=joint_transform_test,
transform=input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=1,
num_workers=4,
shuffle=False)
criterion = CrossEntropyLoss2d(size_average=False,
ignore_index=VOC.ignore_label).cuda()
#optimizer = optim.SGD(net.parameters(), lr = train_args['lr'], momentum=0.9,weight_decay=train_args['weight_decay'])
optimizer = optim.SGD(
[{
'params': [
param for name, param in net.named_parameters()
if name[-4:] == 'bias'
],
'lr':
2 * train_args['lr'],
'momentum':
train_args['momentum'],
'weight_decay':
0
}, {
'params': [
param for name, param in net.named_parameters()
if name[-4:] != 'bias'
],
'lr':
train_args['lr'],
'momentum':
train_args['momentum'],
'weight_decay':
train_args['weight_decay']
}], {
'params': [
param for name, param in net.named_parameters()
if name[-8:] == 'voc.bias'
],
'lr':
20 * train_args['lr'],
'momentum':
train_args['momentum'],
'weight_decay':
0
}, {
'params': [
param for name, param in net.named_parameters()
if name[-10:] != 'voc.weight'
],
'lr':
10 * train_args['lr'],
'momentum':
train_args['momentum'],
'weight_decay':
train_args['weight_decay']
})
if len(train_args['snapshot']) > 0:
optimizer.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name,
'opt_' + train_args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * train_args['lr']
optimizer.param_groups[1]['lr'] = train_args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(
os.path.join(ckpt_path, exp_name,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(train_args) + '\n\n')
#scheduler = ReduceLROnPlateau(optimizer, 'min', patience=train_args['lr_patience'], min_lr=1e-10, verbose=True)
scheduler = StepLR(optimizer, step_size=13, gamma=0.1)
for epoch in range(curr_epoch, train_args['epoch_num'] + 1):
train(train_loader, net, criterion, optimizer, epoch, train_args)
val_loss = validate(val_loader, net, criterion, optimizer, epoch,
train_args, restore_transform, visualize)
#scheduler.step(val_loss)
scheduler.step()
def main():
net = PSPNet(num_classes=cityscapes.num_classes)
if len(args['snapshot']) == 0:
# net.load_state_dict(torch.load(os.path.join(ckpt_path, 'cityscapes (coarse)-psp_net', 'xx.pth')))
curr_epoch = 1
args['best_record'] = {'epoch': 0, 'iter': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0,
'fwavacc': 0}
else:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {'epoch': int(split_snapshot[1]), 'iter': int(split_snapshot[3]),
'val_loss': float(split_snapshot[5]), 'acc': float(split_snapshot[7]),
'acc_cls': float(split_snapshot[9]),'mean_iu': float(split_snapshot[11]),
'fwavacc': float(split_snapshot[13])}
net.cuda().train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(args['longer_size']),
joint_transforms.RandomRotate(10),
joint_transforms.RandomHorizontallyFlip()
])
sliding_crop = joint_transforms.SlidingCrop(args['crop_size'], args['stride_rate'], cityscapes.ignore_label)
train_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
visualize = standard_transforms.Compose([
standard_transforms.Scale(args['val_img_display_size']),
standard_transforms.ToTensor()
])
train_set = cityscapes.CityScapes('fine', 'train', joint_transform=train_joint_transform, sliding_crop=sliding_crop,
transform=train_input_transform, target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=args['train_batch_size'], num_workers=8, shuffle=True)
val_set = cityscapes.CityScapes('fine', 'val', transform=val_input_transform, sliding_crop=sliding_crop,
target_transform=target_transform)
val_loader = DataLoader(val_set, batch_size=1, num_workers=8, shuffle=False)
criterion = CrossEntropyLoss2d(size_average=True, ignore_index=cityscapes.ignore_label).cuda()
optimizer = optim.SGD([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': args['lr'], 'weight_decay': args['weight_decay']}
], momentum=args['momentum'], nesterov=True)
if len(args['snapshot']) > 0:
optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * args['lr']
optimizer.param_groups[1]['lr'] = args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(args) + '\n\n')
train(train_loader, net, criterion, optimizer, curr_epoch, args, val_loader, visualize)
def validate(val_loader, net, criterion, epoch, num_known_classes,
num_unknown_classes, hidden, train_args):
model_list = []
feat_list = []
prds_list = []
true_list = []
# Setting network for evaluation mode.
net.eval()
n_patches = 0
if dataset_name == 'Vaihingen':
n_patches = 907 # Vaihingen.
elif dataset_name == 'Potsdam':
n_patches = 12393 # 8993 # Potsdam.
np.random.seed(12345)
if dataset_name == 'Vaihingen':
perm = np.random.permutation(n_patches)[:150].tolist()
elif dataset_name == 'Potsdam':
perm = np.random.permutation(n_patches)[:100].tolist()
with torch.no_grad():
# Creating output directory.
check_mkdir(os.path.join(outp_path, exp_name, 'epoch_' + str(epoch)))
for i, data in enumerate(val_loader):
print('Validation Batch %d/%d' % (i + 1, len(val_loader)))
sys.stdout.flush()
# Obtaining images, labels and paths for batch.
inps_batch, labs_batch, true_batch, img_name = data
inps_batch = inps_batch.squeeze()
labs_batch = labs_batch.squeeze()
true_batch = true_batch.squeeze()
# Iterating over patches inside batch.
for j in range(inps_batch.size(0)):
for k in range(inps_batch.size(1)):
if count in perm:
inps = inps_batch[j, k].unsqueeze(0)
labs = labs_batch[j, k].unsqueeze(0)
true = true_batch[j, k].unsqueeze(0)
# Casting tensors to cuda.
inps, labs, true = inps.cuda(
args['device']), labs.cuda(
args['device']), true.cuda(args['device'])
# Casting to cuda variables.
inps = Variable(inps).cuda(args['device'])
labs = Variable(labs).cuda(args['device'])
true = Variable(true).cuda(args['device'])
# Forwarding.
if conv_name == 'fcnwideresnet50':
outs, classif1, fv2 = net(inps, feat=True)
elif conv_name == 'fcndensenet121':
outs, classif1, fv2 = net(inps, feat=True)
# Computing loss.
soft_outs = F.softmax(outs, dim=1)
# Obtaining predictions.
prds = soft_outs.data.max(1)[1]
if conv_name == 'fcnwideresnet50':
feat_flat = torch.cat([
outs.squeeze(),
classif1.squeeze(),
fv2.squeeze()
], 0)
elif conv_name == 'fcndensenet121':
feat_flat = torch.cat([
outs.squeeze(),
classif1.squeeze(),
fv2.squeeze()
], 0)
feat_flat = feat_flat.permute(
1, 2, 0).contiguous().view(
feat_flat.size(1) * feat_flat.size(2),
feat_flat.size(0)).cpu().numpy()
prds_flat = prds.cpu().numpy().ravel()
true_flat = true.cpu().numpy().ravel()
feat_list.append(feat_flat)
prds_list.append(prds_flat)
true_list.append(true_flat)
gc.collect()
feat_list = np.asarray(feat_list)
gc.collect()
prds_list = np.asarray(prds_list)
gc.collect()
true_list = np.asarray(true_list)
gc.collect()
feat_list = feat_list.reshape(feat_list.shape[0] * feat_list.shape[1],
feat_list.shape[2])
gc.collect()
prds_list = prds_list.reshape(prds_list.shape[0] * prds_list.shape[1])
gc.collect()
true_list = true_list.reshape(true_list.shape[0] * true_list.shape[1])
gc.collect()
for c in range(num_known_classes):
print('Fitting model for class %d...' % (c))
sys.stdout.flush()
tic = time.time()
# Computing PCA models from features.
model = fit_pca_model(feat_list, true_list, prds_list, c,
args['n_components'])
model_list.append(model)
toc = time.time()
print(' Time spent fitting model %d: %.2f' % (c, toc - tic))
scr_thresholds = fit_quantiles(model_list, feat_list, true_list, prds_list,
c)
model_full = {'generative': model_list, 'thresholds': scr_thresholds}
# Saving model on disk.
model_path = os.path.join(outp_path, exp_name, 'model_pca.pkl')
print('Saving model at "%s"...' % (model_path))
sys.stdout.flush()
joblib.dump(model_full, model_path)
return model_full
def main(train_args):
net = FCN8s(num_classes=voc.num_classes, caffe=True).cuda()
if len(train_args['snapshot']) == 0:
curr_epoch = 1
train_args['best_record'] = {'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0}
else:
print('training resumes from ' + train_args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, train_args['snapshot'])))
split_snapshot = train_args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
train_args['best_record'] = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])}
net.train()
mean_std = ([103.939, 116.779, 123.68], [1.0, 1.0, 1.0])
input_transform = standard_transforms.Compose([
extended_transforms.FlipChannels(),
standard_transforms.ToTensor(),
standard_transforms.Lambda(lambda x: x.mul_(255)),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.Lambda(lambda x: x.div_(255)),
standard_transforms.ToPILImage(),
extended_transforms.FlipChannels()
])
visualize = standard_transforms.Compose([
standard_transforms.Scale(400),
standard_transforms.CenterCrop(400),
standard_transforms.ToTensor()
])
train_set = voc.VOC('train', transform=input_transform, target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=1, num_workers=4, shuffle=True)
val_set = voc.VOC('val', transform=input_transform, target_transform=target_transform)
val_loader = DataLoader(val_set, batch_size=1, num_workers=4, shuffle=False)
criterion = CrossEntropyLoss2d(size_average=False, ignore_index=voc.ignore_label).cuda()
optimizer = optim.SGD([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * train_args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': train_args['lr'], 'weight_decay': train_args['weight_decay']}
], momentum=train_args['momentum'])
if len(train_args['snapshot']) > 0:
optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + train_args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * train_args['lr']
optimizer.param_groups[1]['lr'] = train_args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(train_args) + '\n\n')
scheduler = ReduceLROnPlateau(optimizer, 'min', patience=train_args['lr_patience'], min_lr=1e-10, verbose=True)
for epoch in range(curr_epoch, train_args['epoch_num'] + 1):
train(train_loader, net, criterion, optimizer, epoch, train_args)
val_loss = validate(val_loader, net, criterion, optimizer, epoch, train_args, restore_transform, visualize)
scheduler.step(val_loss)
def test(test_loader, net, criterion, epoch, num_known_classes,
num_unknown_classes, hidden, train_args, save_images, save_model,
model_full):
# Setting network for evaluation mode.
net.eval()
with torch.no_grad():
# Creating output directory.
check_mkdir(os.path.join(outp_path, exp_name, 'epoch_' + str(epoch)))
# Iterating over batches.
for i, data in enumerate(test_loader):
print('Test Batch %d/%d' % (i + 1, len(test_loader)))
sys.stdout.flush()
# Obtaining images, labels and paths for batch.
inps_batch, labs_batch, true_batch, img_name = data
inps_batch = inps_batch.squeeze()
labs_batch = labs_batch.squeeze()
true_batch = true_batch.squeeze()
# Iterating over patches inside batch.
for j in range(inps_batch.size(0)):
print(' Test MiniBatch %d/%d' % (j + 1, inps_batch.size(0)))
sys.stdout.flush()
tic = time.time()
for k in range(inps_batch.size(1)):
inps = inps_batch[j, k].unsqueeze(0)
labs = labs_batch[j, k].unsqueeze(0)
true = true_batch[j, k].unsqueeze(0)
# Casting tensors to cuda.
inps, labs, true = inps.cuda(args['device']), labs.cuda(
args['device']), true.cuda(args['device'])
# Casting to cuda variables.
inps = Variable(inps).cuda(args['device'])
labs = Variable(labs).cuda(args['device'])
true = Variable(true).cuda(args['device'])
# Forwarding.
if conv_name == 'fcnwideresnet50':
outs, classif1, fv2 = net(inps, feat=True)
elif conv_name == 'fcndensenet121':
outs, classif1, fv2 = net(inps, feat=True)
# Computing probabilities.
soft_outs = F.softmax(outs, dim=1)
# Obtaining prior predictions.
prds = soft_outs.data.max(1)[1]
# Obtaining posterior predictions.
if conv_name == 'fcnwideresnet50':
feat_flat = torch.cat([outs, classif1, fv2], 1)
elif conv_name == 'fcndensenet121':
feat_flat = torch.cat([outs, classif1, fv2], 1)
feat_flat = feat_flat.permute(
0, 2, 3, 1).contiguous().view(
feat_flat.size(0) *
feat_flat.size(2) * feat_flat.size(3),
feat_flat.size(1)).cpu().numpy()
prds_flat = prds.cpu().numpy().ravel()
true_flat = true.cpu().numpy().ravel()
prds_post, scores = pred_pixelwise(
model_full, feat_flat, prds_flat, num_known_classes,
model_full['thresholds'][-6])
prds_post = prds_post.reshape(prds.size(0), prds.size(1),
prds.size(2))
scores = scores.reshape(prds.size(0), prds.size(1),
prds.size(2))
# Saving predictions.
if (save_images):
pred_prev_path = os.path.join(
outp_path, exp_name, 'epoch_' + str(epoch),
img_name[0].replace(
'.tif',
'_prev_' + str(j) + '_' + str(k) + '.png'))
scor_path = os.path.join(
outp_path, exp_name, 'epoch_' + str(epoch),
img_name[0].replace(
'.tif',
'_scor_' + str(j) + '_' + str(k) + '.npy'))
io.imsave(
pred_prev_path,
util.img_as_ubyte(prds.cpu().squeeze().numpy()))
np.save(scor_path, scores.squeeze())
toc = time.time()
print(' Elapsed Time: %.2f' % (toc - tic))
sys.stdout.flush()
from os.path import join
from utils import check_mkdir
# Most likely you will need to adjust caffe_root, data_root and *.caffemodel -
# - model_weights in models tuples
# GPU mode is recommended, CPU mode may take several
# days to complete feature extraction
#
caffe_mode = "CPU" # CPU|GPU
caffe_device = 0
caffe_root = "/home/torch/caffe-rc3"
data_root = "/media/torch/WD/kaggle/Yelp"
check_mkdir(join(data_root, "features"))
check_mkdir(join(data_root, "submission"))
#
# models: tuple of (model_weights, model_prototxt, mean_image,
# image_feature_file, biz_feature_file)
#
models = (
(join(caffe_root, "models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel"),
join(data_root, "models/reference_caffenet.prototxt"),
join(data_root, "models/ilsvrc_2012_mean.npy"),
join(data_root, "features/bvlc_reference_{}_image_features.h5"),
join(data_root, "features/bvlc_reference_{}_biz_features_{}.pkl")),
(join(caffe_root, "models/PlacesCNN/places205CNN_iter_300000.caffemodel"),
join(data_root, "models/places205CNN.prototxt"),
def main(train_args):
# Setting network architecture.
if (conv_name == 'fcnwideresnet50'):
net = FCNWideResNet50(4,
num_classes=args['num_classes'],
pretrained=False,
skip=True,
hidden_classes=hidden).cuda(args['device'])
elif (conv_name == 'fcndensenet121'):
net = FCNDenseNet121(4,
num_classes=args['num_classes'],
pretrained=False,
skip=True,
hidden_classes=hidden).cuda(args['device'])
print('Loading pretrained weights from file "' + pretrained_path + '"')
net.load_state_dict(torch.load(pretrained_path))
# print(net)
train_args['best_record'] = {
'epoch': 0,
'lr': 1e-4,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'iou': 0
}
# Setting datasets.
val_set = list_dataset.ListDataset(
dataset_name,
'Validate', (train_args['h_size'], train_args['w_size']),
'statistical',
hidden,
overlap=False,
use_dsm=True)
val_loader = DataLoader(val_set,
batch_size=1,
num_workers=train_args['num_workers'],
shuffle=False)
test_set = list_dataset.ListDataset(
dataset_name,
'Test', (train_args['h_size'], train_args['w_size']),
'statistical',
hidden,
overlap=True,
use_dsm=True)
test_loader = DataLoader(test_set,
batch_size=1,
num_workers=train_args['num_workers'],
shuffle=False)
# Setting criterion.
criterion = CrossEntropyLoss2d(size_average=False,
ignore_index=5).cuda(args['device'])
# Making sure checkpoint and output directories are created.
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
check_mkdir(outp_path)
check_mkdir(os.path.join(outp_path, exp_name))
# Validation function.
model_full = validate(val_loader, net, criterion, epoch, num_known_classes,
num_unknown_classes, hidden, train_args)
# Computing test.
test(test_loader, net, criterion, epoch, num_known_classes,
num_unknown_classes, hidden, train_args, True,
epoch % args['save_freq'] == 0, model_full)
print('Exiting...')
def main():
net = PSPNet(num_classes=voc.num_classes).cuda()
if len(args['snapshot']) == 0:
net.load_state_dict(torch.load(os.path.join(ckpt_path, 'cityscapes (coarse)-psp_net', 'xx.pth')))
curr_epoch = 1
args['best_record'] = {'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0}
else:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])}
net.train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(args['longer_size']),
joint_transforms.RandomRotate(10),
joint_transforms.RandomHorizontallyFlip()
])
sliding_crop = joint_transforms.SlidingCrop(args['crop_size'], args['stride_rate'], voc.ignore_label)
train_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
val_input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
visualize = standard_transforms.Compose([
standard_transforms.Scale(args['val_img_display_size']),
standard_transforms.ToTensor()
])
train_set = voc.VOC('train', joint_transform=train_joint_transform, sliding_crop=sliding_crop,
transform=train_input_transform, target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=args['train_batch_size'], num_workers=8, shuffle=True)
val_set = voc.VOC('val', transform=val_input_transform, sliding_crop=sliding_crop,
target_transform=target_transform)
val_loader = DataLoader(val_set, batch_size=1, num_workers=8, shuffle=False)
criterion = CrossEntropyLoss2d(size_average=True, ignore_index=voc.ignore_label).cuda()
optimizer = optim.SGD([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': args['lr'], 'weight_decay': args['weight_decay']}
], momentum=args['momentum'], nesterov=True)
if len(args['snapshot']) > 0:
optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * args['lr']
optimizer.param_groups[1]['lr'] = args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(args) + '\n\n')
train(train_loader, net, criterion, optimizer, curr_epoch, args, val_loader, visualize)
def main():
net = FCN8s(num_classes=cityscapes.num_classes).cuda()
if len(args['snapshot']) == 0:
curr_epoch = 1
args['best_record'] = {'epoch': 0, 'val_loss': 1e10, 'acc': 0, 'acc_cls': 0, 'mean_iu': 0, 'fwavacc': 0}
else:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {'epoch': int(split_snapshot[1]), 'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]), 'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]), 'fwavacc': float(split_snapshot[11])}
net.train()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
short_size = int(min(args['input_size']) / 0.875)
train_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(short_size),
joint_transforms.RandomCrop(args['input_size']),
joint_transforms.RandomHorizontallyFlip()
])
val_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(short_size),
joint_transforms.CenterCrop(args['input_size'])
])
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage()
])
visualize = standard_transforms.ToTensor()
train_set = cityscapes.CityScapes('fine', 'train', joint_transform=train_joint_transform,
transform=input_transform, target_transform=target_transform)
train_loader = DataLoader(train_set, batch_size=args['train_batch_size'], num_workers=8, shuffle=True)
val_set = cityscapes.CityScapes('fine', 'val', joint_transform=val_joint_transform, transform=input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set, batch_size=args['val_batch_size'], num_workers=8, shuffle=False)
criterion = CrossEntropyLoss2d(size_average=False, ignore_index=cityscapes.ignore_label).cuda()
optimizer = optim.SGD([
{'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
'lr': 2 * args['lr']},
{'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
'lr': args['lr'], 'weight_decay': args['weight_decay']}
], momentum=args['momentum'])
if len(args['snapshot']) > 0:
optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, 'opt_' + args['snapshot'])))
optimizer.param_groups[0]['lr'] = 2 * args['lr']
optimizer.param_groups[1]['lr'] = args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt'), 'w').write(str(args) + '\n\n')
scheduler = ReduceLROnPlateau(optimizer, 'min', patience=args['lr_patience'], min_lr=1e-10)
for epoch in range(curr_epoch, args['epoch_num'] + 1):
train(train_loader, net, criterion, optimizer, epoch, args)
val_loss = validate(val_loader, net, criterion, optimizer, epoch, args, restore_transform, visualize)
scheduler.step(val_loss)
