class Test:
def __init__(self,
model_path,
config,
bn,
save_path,
save_batch,
cuda=False):
self.bn = bn
self.target = config.all_dataset
self.target.remove(config.dataset)
# load source domain
self.source_set = spacenet.Spacenet(city=config.dataset,
split='test',
img_root=config.img_root)
self.source_loader = DataLoader(self.source_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.save_path = save_path
self.save_batch = save_batch
self.target_set = []
self.target_loader = []
self.target_trainset = []
self.target_trainloader = []
self.config = config
# load other domains
for city in self.target:
#test_img_root = '/home/home1/swarnakr/main/DomainAdaptation/satellite/' + city + '/' + 'test'
#test = spacenet.Spacenet(city=city, split='test', img_root=test_img_root)
self.target_set.append(test)
self.target_loader.append(
DataLoader(test, batch_size=16, shuffle=False, num_workers=2))
#train_img_root = '/home/home1/swarnakr/main/DomainAdaptation/satellite/' + city + '/' + 'train'
#train = spacenet.Spacenet(city=city, split='train', img_root=train_img_root)
self.target_trainset.append(train)
self.target_trainloader.append(
DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path,
map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def get_performance(self, dataloader, trainloader, city):
# change mean and var of bn to adapt to the target domain
if self.bn and city != self.config.dataset:
print('BN Adaptation on' + city)
self.model.train()
# print layer params
if 0:
layr = 0
for h in self.model.modules():
if isinstance(h, nn.Conv2d):
k1 = h.kernel_size[0]
k2 = h.kernel_size[1]
ch = h.out_channels
print(
'L={} & ${} \\times {} \\times {}$ \\\\ \\hline'.
format(layr, k1, k2, ch))
layr += 1
for sample in trainloader:
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
#pdb.set_trace()
batch = self.save_batch
self.model.eval()
self.evaluator.reset()
tbar = tqdm(dataloader, desc='\r')
# save in different directories
if self.bn:
save_path = os.path.join(self.save_path, city + '_bn')
else:
save_path = os.path.join(self.save_path, city)
# evaluate on the test dataset
bn = dict()
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
# save BN params
layr = 0
for h in self.model.modules():
if isinstance(
h, nn.Conv2d
): #this is ok to do, since there is one conv2d layer corresponding to each BN layer.
bn[(layr), 'weight'] = np.squeeze(h.weight)
if isinstance(h, nn.BatchNorm2d): #Conv2d):
bn[(layr, 'mean')] = h.running_mean
bn[(layr, 'var')] = h.running_var
layr += 1
#pdb.set_trace()
if not os.path.exists(self.save_path):
os.mkdir(self.save_path)
torch.save(bn, os.path.join(save_path, 'bnAll.pth'))
break
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# save pictures
if batch > 0:
if not os.path.exists(self.save_path):
os.mkdir(self.save_path)
if not os.path.exists(save_path):
os.mkdir(save_path)
image = image.cpu().numpy() * 255
image = image.transpose(0, 2, 3, 1).astype(int)
imgs = self.color_images(pred, target)
self.save_images(imgs, batch, save_path, False)
self.save_images(image, batch, save_path, True)
batch -= 1
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
return Acc, IoU, mIoU
def test(self):
# A=[]; I=[]; Im=[];
#A, I, Im = self.get_performance(self.source_loader, None, self.config.dataset)
tA, tI, tIm = [], [], []
for dl, tl, city in zip(self.target_loader, self.target_trainloader,
self.target):
#if city != 'Vegas':
tA_, tI_, tIm_ = self.get_performance(dl, tl, city)
tA.append(tA_)
tI.append(tI_)
tIm.append(tIm_)
res = {}
print("Test for source domain:")
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(self.config.dataset, A, I,
Im))
res[config.dataset] = {'Acc': A, 'IoU': I, 'mIoU': Im}
print('Test for target domain:')
for i, city in enumerate(self.target):
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(city, tA[i], tI[i],
tIm[i]))
res[city] = {'Acc': tA[i], 'IoU': tI[i], 'mIoU': tIm[i]}
if self.bn:
name = 'train_log/test_bn.json'
else:
name = 'train_log/test.json'
with open(name, 'w') as f:
json.dump(res, f)
def save_images(self, imgs, batch_index, save_path, if_original=False):
for i, img in enumerate(imgs):
img = img[:, :, ::-1] # change to BGR
#from IPython import embed
#embed()
if not if_original:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Original.jpg'),
img)
else:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Pred.jpg'), img)
def color_images(self, pred, target):
imgs = []
for p, t in zip(pred, target):
tmp = p * 2 + t
np.squeeze(tmp)
img = np.zeros((p.shape[0], p.shape[1], 3))
# bkg:negative, building:postive
#from IPython import embed
#embed()
img[np.where(tmp == 0)] = [0, 0, 0] # Black RGB, for true negative
img[np.where(tmp == 1)] = [255, 0,
0] # Red RGB, for false negative
img[np.where(tmp == 2)] = [0, 255,
0] # Green RGB, for false positive
img[np.where(tmp == 3)] = [255, 255,
0] #Yellow RGB, for true positive
imgs.append(img)
return imgs
class Trainer(object):
def __init__(self, config, args):
self.args = args
self.config = config
self.visdom = args.visdom
if args.visdom:
self.vis = visdom.Visdom(env=os.getcwd().split('/')[-1], port=8888)
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
config)
# Define network
self.model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': config.lr
}, {
'params': self.model.get_10x_lr_params(),
'lr': config.lr * 10
}]
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=config.momentum,
weight_decay=config.weight_decay)
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(
weight=None, cuda=args.cuda).build_loss(mode=config.loss)
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
self.summary = TensorboardSummary('train_log')
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.best_pred_source = 0.0
# Resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, sample in enumerate(tbar):
itr = epoch * len(self.train_loader) + i
if self.visdom:
self.vis.line(
X=torch.tensor([itr]),
Y=torch.tensor([self.optimizer.param_groups[0]['lr']]),
win='lr',
opts=dict(title='lr', xlabel='iter', ylabel='lr'),
update='append' if itr > 0 else None)
A_image, A_target = sample['image'], sample['label']
if self.args.cuda:
A_image, A_target = A_image.cuda(), A_target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred_source, 0)
A_output, A_feat, A_low_feat = self.model(A_image)
self.optimizer.zero_grad()
# Supervised loss
seg_loss = self.criterion(A_output, A_target)
loss = seg_loss
loss.backward()
self.optimizer.step()
train_loss += seg_loss.item()
self.summary.writer.add_scalar('Train/Loss', loss.item(), itr)
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + A_image.data.shape[0]))
print('Seg Loss: %.3f' % train_loss)
if self.visdom:
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([seg_loss_sum]),
win='train_loss',
name='Seg_loss',
opts=dict(title='loss',
xlabel='epoch',
ylabel='loss'),
update='append' if epoch > 0 else None)
def validation(self, epoch):
def get_metrics(tbar, if_source=False):
self.evaluator.reset()
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output, low_feat, feat = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
self.summary.writer.add_scalar('Val/Loss', test_loss / (i + 1),
epoch)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
if if_source:
print('Validation on source:')
else:
print('Validation on target:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + image.data.shape[0]))
print("Acc:{}, IoU:{}, mIoU:{}".format(Acc, IoU, mIoU))
print('Loss: %.3f' % test_loss)
if if_source:
names = ['source', 'source_acc', 'source_IoU', 'source_mIoU']
self.summary.writer.add_scalar('Val/SourceAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/SourceIoU', IoU, epoch)
else:
names = ['target', 'target_acc', 'target_IoU', 'target_mIoU']
self.summary.writer.add_scalar('Val/TargetAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/TargetIoU', IoU, epoch)
# Draw Visdom
if self.visdom:
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([test_loss]),
win='val_loss',
name=names[0],
update='append')
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([Acc]),
win='metrics',
name=names[1],
opts=dict(title='metrics',
xlabel='epoch',
ylabel='performance'),
update='append' if epoch > 0 else None)
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([IoU]),
win='metrics',
name=names[2],
update='append')
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([mIoU]),
win='metrics',
name=names[3],
update='append')
return Acc, IoU, mIoU
self.model.eval()
tbar_source = tqdm(self.val_loader, desc='\r')
s_acc, s_iou, s_miou = get_metrics(tbar_source, True)
new_pred_source = s_iou
if new_pred_source > self.best_pred_source:
is_best = True
self.best_pred_source = max(new_pred_source, self.best_pred_source)
print('Saving state, epoch:', epoch)
torch.save(
self.model.module.state_dict(), self.args.save_folder +
'models/' + 'epoch' + str(epoch) + '.pth')
loss_file = {'s_Acc': s_acc, 's_IoU': s_iou, 's_mIoU': s_miou}
with open(
os.path.join(self.args.save_folder, 'eval',
'epoch' + str(epoch) + '.json'), 'w') as f:
json.dump(loss_file, f)
class Trainer(object):
def __init__(self, config, args):
self.args = args
self.config = config
self.visdom = args.visdom
if args.visdom:
self.vis = visdom.Visdom(env=os.getcwd().split('/')[-1], port=8888)
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
config)
self.target_train_loader, self.target_val_loader, self.target_test_loader, _ = make_target_data_loader(
config)
# Define network
self.model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
self.D = Discriminator(num_classes=self.nclass, ndf=16)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': config.lr
}, {
'params': self.model.get_10x_lr_params(),
'lr': config.lr * config.lr_ratio
}]
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.D_optimizer = torch.optim.Adam(self.D.parameters(),
lr=config.lr,
betas=(0.9, 0.99))
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(
weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.entropy_mini_loss = MinimizeEntropyLoss()
self.bottleneck_loss = BottleneckLoss()
self.instance_loss = InstanceLoss()
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
self.summary = TensorboardSummary('./train_log')
# labels for adversarial training
self.source_label = 0
self.target_label = 1
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.D = torch.nn.DataParallel(self.D)
patch_replication_callback(self.D)
self.D = self.D.cuda()
self.best_pred_source = 0.0
self.best_pred_target = 0.0
# Resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
def training(self, epoch):
train_loss, seg_loss_sum, bn_loss_sum, entropy_loss_sum, adv_loss_sum, d_loss_sum, ins_loss_sum = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
self.model.train()
if config.freeze_bn:
self.model.module.freeze_bn()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
target_train_iterator = iter(self.target_train_loader)
for i, sample in enumerate(tbar):
itr = epoch * len(self.train_loader) + i
#if self.visdom:
# self.vis.line(X=torch.tensor([itr]), Y=torch.tensor([self.optimizer.param_groups[0]['lr']]),
# win='lr', opts=dict(title='lr', xlabel='iter', ylabel='lr'),
# update='append' if itr>0 else None)
self.summary.writer.add_scalar(
'Train/lr', self.optimizer.param_groups[0]['lr'], itr)
A_image, A_target = sample['image'], sample['label']
# Get one batch from target domain
try:
target_sample = next(target_train_iterator)
except StopIteration:
target_train_iterator = iter(self.target_train_loader)
target_sample = next(target_train_iterator)
B_image, B_target, B_image_pair = target_sample[
'image'], target_sample['label'], target_sample['image_pair']
if self.args.cuda:
A_image, A_target = A_image.cuda(), A_target.cuda()
B_image, B_target, B_image_pair = B_image.cuda(
), B_target.cuda(), B_image_pair.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred_source,
self.best_pred_target, self.config.lr_ratio)
self.scheduler(self.D_optimizer, i, epoch, self.best_pred_source,
self.best_pred_target, self.config.lr_ratio)
A_output, A_feat, A_low_feat = self.model(A_image)
B_output, B_feat, B_low_feat = self.model(B_image)
#B_output_pair, B_feat_pair, B_low_feat_pair = self.model(B_image_pair)
#B_output_pair, B_feat_pair, B_low_feat_pair = flip(B_output_pair, dim=-1), flip(B_feat_pair, dim=-1), flip(B_low_feat_pair, dim=-1)
self.optimizer.zero_grad()
self.D_optimizer.zero_grad()
# Train seg network
for param in self.D.parameters():
param.requires_grad = False
# Supervised loss
seg_loss = self.criterion(A_output, A_target)
main_loss = seg_loss
# Unsupervised loss
#ins_loss = 0.01 * self.instance_loss(B_output, B_output_pair)
#main_loss += ins_loss
# Train adversarial loss
D_out = self.D(prob_2_entropy(F.softmax(B_output)))
adv_loss = bce_loss(D_out, self.source_label)
main_loss += self.config.lambda_adv * adv_loss
main_loss.backward()
# Train discriminator
for param in self.D.parameters():
param.requires_grad = True
A_output_detach = A_output.detach()
B_output_detach = B_output.detach()
# source
D_source = self.D(prob_2_entropy(F.softmax(A_output_detach)))
source_loss = bce_loss(D_source, self.source_label)
source_loss = source_loss / 2
# target
D_target = self.D(prob_2_entropy(F.softmax(B_output_detach)))
target_loss = bce_loss(D_target, self.target_label)
target_loss = target_loss / 2
d_loss = source_loss + target_loss
d_loss.backward()
self.optimizer.step()
self.D_optimizer.step()
seg_loss_sum += seg_loss.item()
#ins_loss_sum += ins_loss.item()
adv_loss_sum += self.config.lambda_adv * adv_loss.item()
d_loss_sum += d_loss.item()
#train_loss += seg_loss.item() + self.config.lambda_adv * adv_loss.item()
train_loss += seg_loss.item()
self.summary.writer.add_scalar('Train/SegLoss', seg_loss.item(),
itr)
#self.summary.writer.add_scalar('Train/InsLoss', ins_loss.item(), itr)
self.summary.writer.add_scalar('Train/AdvLoss', adv_loss.item(),
itr)
self.summary.writer.add_scalar('Train/DiscriminatorLoss',
d_loss.item(), itr)
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
# Show the results of the last iteration
#if i == len(self.train_loader)-1:
print("Add Train images at epoch" + str(epoch))
self.summary.visualize_image('Train-Source', self.config.dataset,
A_image, A_target, A_output, epoch, 5)
self.summary.visualize_image('Train-Target', self.config.target,
B_image, B_target, B_output, epoch, 5)
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + A_image.data.shape[0]))
print('Loss: %.3f' % train_loss)
#print('Seg Loss: %.3f' % seg_loss_sum)
#print('Ins Loss: %.3f' % ins_loss_sum)
#print('BN Loss: %.3f' % bn_loss_sum)
#print('Adv Loss: %.3f' % adv_loss_sum)
#print('Discriminator Loss: %.3f' % d_loss_sum)
#if self.visdom:
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([seg_loss_sum]), win='train_loss', name='Seg_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([ins_loss_sum]), win='train_loss', name='Ins_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([bn_loss_sum]), win='train_loss', name='BN_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([adv_loss_sum]), win='train_loss', name='Adv_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([d_loss_sum]), win='train_loss', name='Dis_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
def validation(self, epoch):
def get_metrics(tbar, if_source=False):
self.evaluator.reset()
test_loss = 0.0
#feat_mean, low_feat_mean, feat_var, low_feat_var = 0, 0, 0, 0
#adv_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output, low_feat, feat = self.model(image)
#low_feat = low_feat.cpu().numpy()
#feat = feat.cpu().numpy()
#if isinstance(feat, np.ndarray):
# feat_mean += feat.mean(axis=0).mean(axis=1).mean(axis=1)
# low_feat_mean += low_feat.mean(axis=0).mean(axis=1).mean(axis=1)
# feat_var += feat.var(axis=0).var(axis=1).var(axis=1)
# low_feat_var += low_feat.var(axis=0).var(axis=1).var(axis=1)
#else:
# feat_mean = feat.mean(axis=0).mean(axis=1).mean(axis=1)
# low_feat_mean = low_feat.mean(axis=0).mean(axis=1).mean(axis=1)
# feat_var = feat.var(axis=0).var(axis=1).var(axis=1)
# low_feat_var = low_feat.var(axis=0).var(axis=1).var(axis=1)
#d_output = self.D(prob_2_entropy(F.softmax(output)))
#adv_loss += bce_loss(d_output, self.source_label).item()
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target_ = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target_, pred)
if if_source:
print("Add Validation-Source images at epoch" + str(epoch))
self.summary.visualize_image('Val-Source', self.config.dataset,
image, target, output, epoch, 5)
else:
print("Add Validation-Target images at epoch" + str(epoch))
self.summary.visualize_image('Val-Target', self.config.target,
image, target, output, epoch, 5)
#feat_mean /= (i+1)
#low_feat_mean /= (i+1)
#feat_var /= (i+1)
#low_feat_var /= (i+1)
#adv_loss /= (i+1)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
if if_source:
print('Validation on source:')
else:
print('Validation on target:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + image.data.shape[0]))
print("Acc:{}, IoU:{}, mIoU:{}".format(Acc, IoU, mIoU))
print('Loss: %.3f' % test_loss)
if if_source:
names = ['source', 'source_acc', 'source_IoU', 'source_mIoU']
self.summary.writer.add_scalar('Val/SourceAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/SourceIoU', IoU, epoch)
else:
names = ['target', 'target_acc', 'target_IoU', 'target_mIoU']
self.summary.writer.add_scalar('Val/TargetAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/TargetIoU', IoU, epoch)
# Draw Visdom
#if if_source:
# names = ['source', 'source_acc', 'source_IoU', 'source_mIoU']
#else:
# names = ['target', 'target_acc', 'target_IoU', 'target_mIoU']
#if self.visdom:
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([test_loss]), win='val_loss', name=names[0],
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([adv_loss]), win='val_loss', name='adv_loss',
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([Acc]), win='metrics', name=names[1],
# opts=dict(title='metrics', xlabel='epoch', ylabel='performance'),
# update='append' if epoch > 0 else None)
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([IoU]), win='metrics', name=names[2],
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([mIoU]), win='metrics', name=names[3],
# update='append')
return Acc, IoU, mIoU
self.model.eval()
tbar_source = tqdm(self.val_loader, desc='\r')
tbar_target = tqdm(self.target_val_loader, desc='\r')
s_acc, s_iou, s_miou = get_metrics(tbar_source, True)
t_acc, t_iou, t_miou = get_metrics(tbar_target, False)
new_pred_source = s_iou
new_pred_target = t_iou
if new_pred_source > self.best_pred_source or new_pred_target > self.best_pred_target:
is_best = True
self.best_pred_source = max(new_pred_source, self.best_pred_source)
self.best_pred_target = max(new_pred_target, self.best_pred_target)
print('Saving state, epoch:', epoch)
torch.save(
self.model.module.state_dict(),
self.args.save_folder + 'models/' + 'epoch' + str(epoch) + '.pth')
loss_file = {
's_Acc': s_acc,
's_IoU': s_iou,
's_mIoU': s_miou,
't_Acc': t_acc,
't_IoU': t_iou,
't_mIoU': t_miou
}
with open(
os.path.join(self.args.save_folder, 'eval',
'epoch' + str(epoch) + '.json'), 'w') as f:
json.dump(loss_file, f)
class Test:
def __init__(self,
model_path,
config,
bn,
save_path,
save_batch,
cuda=False):
self.bn = bn
self.target = config.all_dataset
self.target.remove(config.dataset)
# load source domain
self.source_set = spacenet.Spacenet(city=config.dataset,
split='val',
img_root=config.img_root)
self.source_loader = DataLoader(self.source_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.save_path = save_path
self.save_batch = save_batch
self.target_set = []
self.target_loader = []
self.target_trainset = []
self.target_trainloader = []
self.config = config
# load other domains
for city in self.target:
test = spacenet.Spacenet(city=city,
split='val',
img_root=config.img_root)
self.target_set.append(test)
self.target_loader.append(
DataLoader(test, batch_size=16, shuffle=False, num_workers=2))
train = spacenet.Spacenet(city=city,
split='train',
img_root=config.img_root)
self.target_trainset.append(train)
self.target_trainloader.append(
DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path,
map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def save_output(module, input, output):
global activation, i
# save output
print('I came here')
channels = output.permute(1, 0, 2, 3)
c = channels.shape[0]
features = channels.reshape(c, -1)
if len(activation) == i:
activation.append(features)
else:
activation[i] = torch.cat([activation[i], features], dim=1)
i += 1
return
def get_performance(self, dataloader, trainloader, city):
if 1:
pix = torch.load("pix" + self.config.dataset + "_" + city + ".pt")
for k in range(0, len(pix)):
fig = plt.figure()
plt.hist(pix[k])
plt.xlabel('Activation values')
plt.ylabel("Count")
#plt.legend()
fig.savefig('./train_log/figs/pix_' + self.config.dataset +
'_' + city + 'act' + str(k) + '.png')
return [], [], []
# change mean and var of bn to adapt to the target domain
if self.bn and city != self.config.dataset:
self.checkpoint = torch.load('./train_log/' + self.config.dataset +
'_da_' + city + '.pth')
self.model.load_state_dict(self.checkpoint)
if self.cuda:
self.model = self.model.cuda()
if 0: #self.bn and city != self.config.dataset:
print('BN Adaptation on' + city)
self.model.train()
tbar = tqdm(dataloader, desc='\r')
for sample in trainloader:
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
batch = self.save_batch
self.model.eval()
self.evaluator.reset()
tbar = tqdm(dataloader, desc='\r')
# save in different directories
if self.bn:
save_path = os.path.join(self.save_path, city + '_bn')
else:
save_path = os.path.join(self.save_path, city)
global randCh
#randCh={};
global first
first = 1
global ii
ii = 0
global ncity
ncity = city
randCh = torch.load('randCh.pt')
#if city != self.config.dataset:
#randCh = torch.load('randCh.pt')
#else:
# randCh={};
layr = 0
aa = 0
if city == self.config.dataset:
for hh in self.model.modules():
if isinstance(hh, nn.ReLU6): #Conv2d):
layr += 1
if layr % 5 == 0:
#if first==1 and city == self.config.dataset:
#randCh[aa] = np.random.randint(hh.out_channels)
#print(layr)
hh.register_forward_hook(save_output2)
aa += 1
# evaluate on the test dataset
pix = {}
pix1 = {}
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
#manipulate activations here
for k in activation.keys():
if first == 1:
pix[k] = []
pix1[k] = []
for row in range(0, activation[k].shape[1]):
actkrow = activation[k][:, row, :-1].reshape(
-1).cpu().numpy()
pix[k] = np.hstack((pix[k], actkrow))
actkrow1 = activation[k][:, row,
1:].reshape(-1).cpu().numpy()
pix1[k] = np.hstack((pix1[k], actkrow1))
for bb in range(0, activation[k].size(0)):
cv2.imwrite(
os.path.join(
save_path, 'act' + str(k) + 'im' + str(i) +
'b' + str(bb) + city + '.jpg'),
activation[k][bb, :].cpu().numpy() * 255)
first += 1
ii = 0
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# save pictures
if batch > 0:
if not os.path.exists(self.save_path):
os.mkdir(self.save_path)
if not os.path.exists(save_path):
os.mkdir(save_path)
image = image.cpu().numpy() * 255
image = image.transpose(0, 2, 3, 1).astype(int)
imgs = self.color_images(pred, target)
self.save_images(imgs, batch, save_path, False)
self.save_images(image, batch, save_path, True)
batch -= 1
corrVal = {}
for k in activation.keys():
corrVal[k] = np.corrcoef(pix[k], pix1[k])[0, 1]
#_ = plt.hist(pix[k]); plt.show()
torch.save(pix, "pix" + self.config.dataset + "_" + city + ".pt")
print(corrVal)
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
return Acc, IoU, mIoU
def test(self):
A, I, Im = self.get_performance(self.source_loader, None,
self.config.dataset)
tA, tI, tIm = [], [], []
for dl, tl, city in zip(self.target_loader, self.target_trainloader,
self.target):
tA_, tI_, tIm_ = self.get_performance(dl, tl, city)
tA.append(tA_)
tI.append(tI_)
tIm.append(tIm_)
res = {}
print("Test for source domain:")
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(self.config.dataset, A, I,
Im))
res[config.dataset] = {'Acc': A, 'IoU': I, 'mIoU': Im}
print('Test for target domain:')
for i, city in enumerate(self.target):
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(city, tA[i], tI[i],
tIm[i]))
res[city] = {'Acc': tA[i], 'IoU': tI[i], 'mIoU': tIm[i]}
if self.bn:
name = 'train_log/test_bn.json'
else:
name = 'train_log/test.json'
with open(name, 'w') as f:
json.dump(res, f)
def save_images(self, imgs, batch_index, save_path, if_original=False):
for i, img in enumerate(imgs):
img = img[:, :, ::-1] # change to BGR
#from IPython import embed
#embed()
if not if_original:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Original.jpg'),
img)
else:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Pred.jpg'), img)
def color_images(self, pred, target):
imgs = []
for p, t in zip(pred, target):
tmp = p * 2 + t
np.squeeze(tmp)
img = np.zeros((p.shape[0], p.shape[1], 3))
# bkg:negative, building:postive
#from IPython import embed
#embed()
img[np.where(tmp == 0)] = [0, 0, 0] # Black RGB, for true negative
img[np.where(tmp == 1)] = [255, 0,
0] # Red RGB, for false negative
img[np.where(tmp == 2)] = [0, 255,
0] # Green RGB, for false positive
img[np.where(tmp == 3)] = [255, 255,
0] #Yellow RGB, for true positive
imgs.append(img)
return imgs
class Test:
def __init__(self,
model_path,
config,
bn,
save_path,
save_batch,
cuda=False):
self.bn = bn
self.target = config.all_dataset
self.target.remove(config.dataset)
# load source domain
self.source_set = spacenet.Spacenet(city=config.dataset,
split='test',
img_root=config.img_root,
source_dist=dist[config.dataset])
self.source_loader = DataLoader(self.source_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.save_path = save_path
self.save_batch = save_batch
self.target_set = []
self.target_loader = []
self.target_trainset = []
self.target_trainloader = []
self.config = config
# load other domains
for city in self.target:
test = spacenet.Spacenet(city=city,
split='test',
img_root=config.img_root,
source_dist=dist[city])
self.target_set.append(test)
self.target_loader.append(
DataLoader(test, batch_size=16, shuffle=False, num_workers=2))
train = spacenet.Spacenet(city=city,
split='train',
img_root=config.img_root,
source_dist=dist[city])
self.target_trainset.append(train)
self.target_trainloader.append(
DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=False)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path,
map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def get_performance(self, dataloader, trainloader, city):
# change mean and var of bn to adapt to the target domain
if self.bn and city != self.config.dataset:
print('BN Adaptation on ' + city)
self.model.train()
for sample in trainloader:
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
if len(output) > 1:
output = output[0]
batch = self.save_batch
if batch < 0:
batch = len(dataloader)
self.model.eval()
self.evaluator.reset()
tbar = tqdm(dataloader, desc='\r')
# save in different directories
if self.bn:
save_path = os.path.join(self.save_path, city + '_bn')
else:
save_path = os.path.join(self.save_path, city)
# evaluate on the test dataset
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
if len(output) > 1:
output = output[0]
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# save pictures
if batch > 0:
if not os.path.exists(self.save_path):
os.mkdir(self.save_path)
if not os.path.exists(save_path):
os.mkdir(save_path)
image = image.cpu().numpy() * 255
image = image.transpose(0, 2, 3, 1).astype(int)
imgs = self.color_images(pred, target)
self.save_images(imgs, batch, save_path, False)
self.save_images(image, batch, save_path, True)
batch -= 1
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
return Acc, IoU, mIoU
def test(self):
A, I, Im = self.get_performance(self.source_loader, None,
self.config.dataset)
tA, tI, tIm = [], [], []
for dl, tl, city in zip(self.target_loader, self.target_trainloader,
self.target):
tA_, tI_, tIm_ = self.get_performance(dl, tl, city)
tA.append(tA_)
tI.append(tI_)
tIm.append(tIm_)
res = {}
print("Test for source domain:")
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(self.config.dataset, A, I,
Im))
res[config.dataset] = {'Acc': A, 'IoU': I, 'mIoU': Im}
print('Test for target domain:')
for i, city in enumerate(self.target):
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(city, tA[i], tI[i],
tIm[i]))
res[city] = {'Acc': tA[i], 'IoU': tI[i], 'mIoU': tIm[i]}
if self.bn:
name = 'train_log/test_bn.json'
else:
name = 'train_log/test.json'
with open(name, 'w') as f:
json.dump(res, f)
def save_images(self, imgs, batch_index, save_path, if_original=False):
for i, img in enumerate(imgs):
img = img[:, :, ::-1] # change to BGR
#from IPython import embed
#embed()
if if_original:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Original.jpg'),
img)
else:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Pred.jpg'), img)
def color_images(self, pred, target):
imgs = []
for p, t in zip(pred, target):
tmp = p * 2 + t
np.squeeze(tmp)
img = np.zeros((p.shape[0], p.shape[1], 3))
# bkg:negative, building:postive
#from IPython import embed
#embed()
img[np.where(tmp == 0)] = [0, 0, 0] # Black RGB, for true negative
img[np.where(tmp == 1)] = [255, 0,
0] # Red RGB, for false negative
img[np.where(tmp == 2)] = [0, 255,
0] # Green RGB, for false positive
img[np.where(tmp == 3)] = [255, 255,
0] #Yellow RGB, for true positive
imgs.append(img)
return imgs
class Trainer(object):
def __init__(self, config, args):
self.args = args
self.config = config
self.vis = visdom.Visdom(env=os.getcwd().split('/')[-1])
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
config)
self.gpu = args.gpu
# labels for adversarial training
self.gt_label = 0
self.prediction_label = 1
self.argmax = ArgMax()
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
model_D = FCDiscriminator(num_classes=self.nclass)
model_D.train()
train_params = [{
'params': model.get_1x_lr_params(),
'lr': config.lr
}, {
'params': model.get_10x_lr_params(),
'lr': config.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=config.momentum,
weight_decay=config.weight_decay)
optimizer_D = torch.optim.Adam(model_D.parameters(),
lr=1e-4,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(
weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.criterion_D = torch.nn.BCEWithLogitsLoss()
self.model, self.optimizer = model, optimizer
self.model_D, self.optimizer_D = model_D, optimizer_D
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
# Using cuda
if args.cuda:
# self.model = torch.nn.DataParallel(self.model)
# patch_replication_callback(self.model)
# self.model_D = torch.nn.DataParallel(self.model)
# patch_replication_callback(self.model_D)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.model_D = self.model_D.cuda()
self.argmax = self.argmax
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
def training(self, epoch):
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
seg_loss = 0.0
adv_loss = 0.0
D_loss = 0.0
self.model.train()
for i, sample in enumerate(tbar):
iter = epoch * len(self.train_loader) + i
self.vis.line(X=torch.tensor([iter]),
Y=torch.tensor(
[self.optimizer.param_groups[0]['lr']]),
win='lr_seg',
opts=dict(title='lr', xlabel='iter', ylabel='lr'),
update='append' if iter > 0 else None)
self.vis.line(X=torch.tensor([iter]),
Y=torch.tensor(
[self.optimizer_D.param_groups[0]['lr']]),
win='lr_adv',
opts=dict(title='lr', xlabel='iter', ylabel='lr'),
update='append' if iter > 0 else None)
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
self.optimizer_D.zero_grad()
# train the segmentation network
# don't accumulate grads in D
for param in self.model_D.parameters():
param.requires_grad = False
# seg loss
output = self.model(image)
loss1 = self.criterion(output, target)
loss1.backward()
self.optimizer.step()
seg_loss += loss1.item()
# adv loss
print(self.argmax.apply(output))
D_out = self.model_D(self.argmax.apply(output))
loss2 = self.criterion_D(
D_out,
Variable(
torch.FloatTensor(D_out.data.size()).fill_(
self.gt_label)).cuda(self.gpu))
loss2.backward()
self.optimizer_D.step()
adv_loss += loss2.item()
# train the discriminator
# bring back requires_grad
for param in self.model_D.parameters():
param.requires_grad = True
# train_with_prediction
output = output.detach()
D_out1 = self.model_D(self.argmax.apply(output))
loss_D1 = self.criterion_D(
D_out1,
Variable(
torch.FloatTensor(D_out1.data.size()).fill_(
self.prediction_label)).cuda(self.gpu))
loss_D1.backward()
D_loss += loss_D1.data.cpu().numpy()
# train with gt
D_out2 = self.model_D(target)
loss_D2 = self.criterion_D(
D_out2,
Variable(
torch.FloatTensor(D_out2.data.size()).fill_(
self.gt_label)).cuda(self.gpu))
loss_D2.backward()
D_loss += loss_D2.data.cpu().numpy()
tbar.set_description('[Train] Seg loss: %.3f, Adv loss: %.3f, D loss: %.3f' \
% (seg_loss / (i + 1), adv_loss / (i + 1), D_loss / (i + 1)))
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + image.data.shape[0]))
print('Seg loss: %.3f, Adv loss: %.3f, D loss: %.3f' \
% (seg_loss / (i + 1), adv_loss / (i + 1), D_loss / (i + 1)))
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([seg_loss]),
win='seg_loss',
name='train',
opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
update='append' if epoch > 0 else None)
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([adv_loss]),
win='adv_loss',
name='train',
opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
update='append' if epoch > 0 else None)
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([D_loss]),
win='D_loss',
name='train',
opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
update='append' if epoch > 0 else None)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
seg_loss = self.criterion(output, target)
seg_loss += seg_loss.item()
output = output.detach()
D_out = self.model_D(self.argmax(output))
adv_loss = self.criterion_D(
D_out,
Variable(
torch.FloatTensor(D_out.data.size()).fill_(
self.gt_label)).cuda(self.gpu))
tbar.set_description('[Test] Seg loss: %.3f, Adv loss: %.3f' %
(seg_loss / (i + 1), adv_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
# Acc_class = self.evaluator.Pixel_Accuracy_Class()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
# FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
print('Validation:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + image.data.shape[0]))
print("Acc:{}, IoU:{}, mIoU:{}".format(Acc, IoU, mIoU))
print('Seg loss: %.3f, Adv loss: %.3f' % (seg_loss /
(i + 1), adv_loss / (i + 1)))
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([seg_loss]),
win='seg_loss',
name='val',
update='append')
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([adv_loss]),
win='adv_loss',
name='val',
update='append')
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([Acc]),
win='metrics',
name='acc',
opts=dict(title='metrics',
xlabel='epoch',
ylabel='performance'),
update='append' if epoch > 0 else None)
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([IoU]),
win='metrics',
name='IoU',
update='append')
self.vis.line(X=torch.tensor([epoch]),
Y=torch.tensor([mIoU]),
win='metrics',
name='mIoU',
update='append')
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
print('Saving state, epoch:', epoch)
torch.save(
self.model.state_dict(), self.args.save_folder + 'models/' +
'epoch' + str(epoch) + '.pth')
loss_file = {'Acc': Acc, 'IoU': IoU, 'mIoU': mIoU}
with open(
os.path.join(self.args.save_folder, 'eval',
'epoch' + str(epoch) + '.json'), 'w') as f:
json.dump(loss_file, f)
class Test:
def __init__(self, model_path, config, cuda=False):
self.target = config.all_dataset
self.target.remove(config.dataset)
# load source domain
self.source_set = spacenet.Spacenet(city=config.dataset,
split='test',
img_root=config.img_root)
self.source_loader = DataLoader(self.source_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.target_set = []
self.target_loader = []
self.target_trainset = []
self.target_trainloader = []
# load other domains
for city in self.target:
test = spacenet.Spacenet(city=city,
split='test',
img_root=config.img_root)
train = spacenet.Spacenet(city=city,
split='train',
img_root=config.img_root)
self.target_set.append(test)
self.target_trainset.append(train)
self.target_loader.append(
DataLoader(test, batch_size=16, shuffle=False, num_workers=2))
self.target_trainloader.append(
DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path,
map_location=torch.device('cpu'))
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def get_performance(self, dataloader, trainloader, if_source):
# change mean and var of bn to adapt to the target domain
if not if_source:
self.model.train()
for sample in trainloader:
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
# evaluate the model
self.model.eval()
self.evaluator.reset()
tbar = tqdm(dataloader, desc='\r')
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
return Acc, IoU, mIoU
def test(self):
A, I, Im = self.get_performance(self.source_loader, None, True)
tA, tI, tIm = [], [], []
for dl, tl in zip(self.target_loader, self.target_trainloader):
tA_, tI_, tIm_ = self.get_performance(dl, tl, False)
tA.append(tA_)
tI.append(tI_)
tIm.append(tIm_)
res = {}
print("Test for source domain:")
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(config.dataset, A, I, Im))
res[config.dataset] = {'Acc': A, 'IoU': I, 'mIoU': Im}
print('Test for target domain:')
for i, city in enumerate(self.target):
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(city, tA[i], tI[i],
tIm[i]))
res[city] = {'Acc': tA[i], 'IoU': tI[i], 'mIoU': tIm[i]}
with open('train_log/test_bn.json', 'w') as f:
json.dump(res, f)
class Test:
def __init__(self, model_path, source, target, cuda=False):
self.source_set = spacenet.Spacenet(city=source,
split='test',
img_root=config.img_root)
self.target_set = spacenet.Spacenet(city=target,
split='test',
img_root=config.img_root)
self.source_loader = DataLoader(self.source_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.target_loader = DataLoader(self.target_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path,
map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def get_performance(self, dataloader):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(dataloader, desc='\r')
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
return Acc, IoU, mIoU
def test(self):
sA, sI, sIm = self.get_performance(self.source_loader)
tA, tI, tIm = self.get_performance(self.target_loader)
print('Test for source domain:')
print("Acc:{}, IoU:{}, mIoU:{}".format(sA, sI, sIm))
print('Test for target domain:')
print("Acc:{}, IoU:{}, mIoU:{}".format(tA, tI, tIm))
res = {
'source': {
'Acc': sA,
'IoU': sI,
'mIoU': sIm
},
'target': {
'Acc': tA,
'IoU': tI,
'mIoU': tIm
}
}
with open('train_log/test.json', 'w') as f:
json.dump(res, f)
class Trainer(object):
def __init__(self, config, args):
self.args = args
self.config = config
#self.vis = visdom.Visdom(env=os.getcwd().split('/')[-1])
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(config)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
train_params = [{'params': model.get_1x_lr_params(), 'lr': config.lr},
{'params': model.get_10x_lr_params(), 'lr': config.lr * 10}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params, momentum=config.momentum,
weight_decay=config.weight_decay)
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.model, self.optimizer = model, optimizer
#pdb.set_trace()
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler, config.lr,
config.epochs, len(self.train_loader),
config.lr_step, config.warmup_epochs)
# Using cuda
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
#self.model = self.model.cuda()
self.model = self.model.to(device)
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
#self.model.module.load_state_dict(checkpoint)
#pdb.set_trace()
if 1:
self.model.module.load_state_dict(checkpoint['model'])
self.args.start_epoch = checkpoint['epoch'] + 1
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['pred']
self.scheduler = checkpoint['scheduler']
#all the above statements would go here
else:
self.model.load_state_dict(checkpoint, map_location=torch.device('cpu'))
#the remaining statements have to re-written
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, args.start_epoch))
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, sample in enumerate(tbar):
iter = epoch * len(self.train_loader) + i
#self.vis.line(X=torch.tensor([iter]), Y=torch.tensor([self.optimizer.param_groups[0]['lr']]),
# win='lr', opts=dict(title='lr', xlabel='iter', ylabel='lr'),
# update='append' if iter>0 else None)
image, target, path = sample['image'], sample['label'], sample['path']
#print(path)
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image)
loss = self.criterion(output, target)
loss.backward()
self.optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.config.batch_size + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([train_loss]), win='loss', name='train',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
# Acc_class = self.evaluator.Pixel_Accuracy_Class()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
# FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
print('Validation:')
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.config.batch_size + image.data.shape[0]))
print("Acc:{}, IoU:{}, mIoU:{}".format(Acc, IoU, mIoU))
print('Loss: %.3f' % test_loss)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([test_loss]), win='loss', name='val',
# update='append')
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([Acc]), win='metrics', name='acc',
# opts=dict(title='metrics', xlabel='epoch', ylabel='performance'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([IoU]), win='metrics', name='IoU',
# update='append')
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([mIoU]), win='metrics', name='mIoU',
# update='append')
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
print('Saving state, epoch:', epoch)
#pdb.set_trace()
state = {
'epoch': epoch,
'optimizer': self.optimizer.state_dict(),
'model': self.model.module.state_dict(),
'pred': new_pred,
'scheduler': self.scheduler
}
torch.save(state, self.args.save_folder + 'models/'
+ 'epoch' + str(epoch) + '.pth')
loss_file = {'Acc': Acc, 'IoU': IoU, 'mIoU': mIoU, 'loss': test_loss}
with open(os.path.join(self.args.save_folder, 'eval', 'epoch' + str(epoch) + '.json'), 'w') as f:
json.dump(loss_file, f)
class Test:
def __init__(self,
model_path,
config,
bn,
save_path,
save_batch,
cuda=False):
self.bn = bn
self.city = config.dataset #all_dataset
self.save_path = save_path
self.save_batch = save_batch
self.target_trainset = []
self.target_trainloader = []
self.config = config
# load other domains
if 1: #for city in self.target:
train = spacenet.Spacenet(city=self.city,
split='train',
img_root=config.img_root)
self.target_trainset.append(train)
self.target_trainloader.append(
DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
#if DA images
self.checkpoint = torch.load(model_path)
#'./train_log/' + self.config.dataset + '_da_' + city + '.pth')
self.model.load_state_dict(self.checkpoint)
if self.cuda:
self.model = self.model.cuda()
def get_performance(self, trainloader, city):
batch = self.save_batch
self.model.eval()
self.evaluator.reset()
tbar = tqdm(trainloader, desc='\r')
# save in different directories
if self.bn:
save_path = os.path.join(self.save_path, city + '_bn')
else:
save_path = os.path.join(self.save_path, city)
if not os.path.exists(self.save_path):
os.mkdir(self.save_path)
layr = 0
for hh in self.model.modules():
if isinstance(hh, nn.ReLU6): #Conv2d):
#hh.register_forward_hook(save_output2)
layr += 1
if layr == 34: #12
hh.register_forward_hook(save_output2)
break
# evaluate on the test dataset
allFeats = {}
for i, sample in enumerate(tbar):
if 1:
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# save pictures
if batch > 0:
if not os.path.exists(self.save_path):
os.mkdir(self.save_path)
if not os.path.exists(save_path):
os.mkdir(save_path)
image = image.cpu().numpy() * 255
image = image.transpose(0, 2, 3, 1).astype(int)
imgs = self.color_images(pred, target)
self.save_images(imgs, batch, save_path, False)
self.save_images(image, batch, save_path, True)
batch -= 1
# pdb.set_trace()
torch.save(activation, save_path + "Feats_" + str(i) + ".pt")
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
return Acc, IoU, mIoU
def test(self):
tA, tI, tIm = [], [], []
for tl, city in zip(self.target_trainloader, self.city):
tA_, tI_, tIm_ = self.get_performance(tl, city)
tA.append(tA_)
tI.append(tI_)
tIm.append(tIm_)
print('Test for target domain:')
for i, city in enumerate(self.city):
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(city, tA[i], tI[i],
tIm[i]))
res[city] = {'Acc': tA[i], 'IoU': tI[i], 'mIoU': tIm[i]}
if self.bn:
name = 'train_log/test_bn.json'
else:
name = 'train_log/test.json'
with open(name, 'w') as f:
json.dump(res, f)
def save_images(self, imgs, batch_index, save_path, if_original=False):
for i, img in enumerate(imgs):
img = img[:, :, ::-1] # change to BGR
#from IPython import embed
#embed()
if not if_original:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Original.jpg'),
img)
else:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Pred.jpg'), img)
def color_images(self, pred, target):
imgs = []
for p, t in zip(pred, target):
tmp = p * 2 + t
np.squeeze(tmp)
img = np.zeros((p.shape[0], p.shape[1], 3))
# bkg:negative, building:postive
#from IPython import embed
#embed()
img[np.where(tmp == 0)] = [0, 0, 0] # Black RGB, for true negative
img[np.where(tmp == 1)] = [255, 0,
0] # Red RGB, for false negative
img[np.where(tmp == 2)] = [0, 255,
0] # Green RGB, for false positive
img[np.where(tmp == 3)] = [255, 255,
0] #Yellow RGB, for true positive
imgs.append(img)
return imgs
class Test:
def __init__(self,
model_path,
config,
bn,
save_path,
save_batch,
cuda=False):
self.bn = bn
self.target = config.all_dataset
self.target.remove(config.dataset)
# load source domain
self.source_set = spacenet.Spacenet(city=config.dataset,
split='test',
img_root=config.img_root)
self.source_loader = DataLoader(self.source_set,
batch_size=16,
shuffle=False,
num_workers=2)
self.save_path = save_path
self.save_batch = save_batch
self.target_set = []
self.target_loader = []
self.target_trainset = []
self.target_trainloader = []
self.config = config
# load other domains
for city in self.target:
test = spacenet.Spacenet(city=city,
split='test',
img_root=config.img_root)
self.target_set.append(test)
self.target_loader.append(
DataLoader(test, batch_size=16, shuffle=False, num_workers=2))
train = spacenet.Spacenet(city=city,
split='train',
img_root=config.img_root)
self.target_trainset.append(train)
self.target_trainloader.append(
DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path,
map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint)
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def save_output(module, input, output):
global activation, i
# save output
print('I came here')
pdb.set_trace()
channels = output.permute(1, 0, 2, 3)
c = channels.shape[0]
features = channels.reshape(c, -1)
if len(activation) == i:
activation.append(features)
else:
activation[i] = torch.cat([activation[i], features], dim=1)
i += 1
return
def get_performance(self, dataloader, trainloader, city):
# change mean and var of bn to adapt to the target domain
#pdb.set_trace()
if self.bn:
save_path = os.path.join(self.save_path, city + '_bn')
else:
save_path = os.path.join(self.save_path, city)
if self.bn and city != self.config.dataset: #!= self.config.dataset:
print('BN Adaptation on' + city)
self.model.train()
layr = 0
for h in self.model.modules():
if isinstance(h, nn.ReLU6): #Conv2d):
layr += 1
if layr == 1:
h.register_forward_hook(save_output2)
if layr > 1:
break
tbar = tqdm(dataloader, desc='\r')
for sample in trainloader:
# for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
pdb.set_trace()
# add0 = np.tile([10, 42, 37],(400,400,16,1));
# add = add0.transpose(2,3,0,1)
# image = image + add;
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
if not os.path.exists(save_path):
os.mkdir(save_path)
pdb.set_trace()
self.save_act(activation, save_path, False)
self.save_act(image.numpy() * 255, save_path, True)
batch = self.save_batch
self.model.eval()
self.evaluator.reset()
tbar = tqdm(dataloader, desc='\r')
# save in different directories
if self.bn:
save_path = os.path.join(self.save_path, city + '_bn')
else:
save_path = os.path.join(self.save_path, city)
layr = 0
for h in self.model.modules():
if isinstance(h, nn.ReLU6): #Conv2d):
layr += 1
if layr == 2:
h.register_forward_hook(save_output2)
if layr > 2:
break
# evaluate on the test dataset
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
pdb.set_trace()
if not os.path.exists(save_path):
os.mkdir(save_path)
self.save_act(activation, save_path, False)
self.save_act(image.numpy() * 255, save_path, True)
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# save pictures
if batch > 0:
if not os.path.exists(self.save_path):
os.mkdir(self.save_path)
if not os.path.exists(save_path):
os.mkdir(save_path)
image = image.cpu().numpy() * 255
image = image.transpose(0, 2, 3, 1).astype(int)
imgs = self.color_images(pred, target)
self.save_images(imgs, batch, save_path, False)
self.save_images(image, batch, save_path, True)
batch -= 1
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
return Acc, IoU, mIoU
def test(self):
#A, I, Im = self.get_performance(self.source_loader, None, self.config.dataset)
tA, tI, tIm = [], [], []
for dl, tl, city in zip(self.target_loader, self.target_trainloader,
self.target):
tA_, tI_, tIm_ = self.get_performance(dl, tl, city)
tA.append(tA_)
tI.append(tI_)
tIm.append(tIm_)
res = {}
print("Test for source domain:")
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(self.config.dataset, A, I,
Im))
res[config.dataset] = {'Acc': A, 'IoU': I, 'mIoU': Im}
print('Test for target domain:')
for i, city in enumerate(self.target):
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(city, tA[i], tI[i],
tIm[i]))
res[city] = {'Acc': tA[i], 'IoU': tI[i], 'mIoU': tIm[i]}
if self.bn:
name = 'train_log/test_bn.json'
else:
name = 'train_log/test.json'
with open(name, 'w') as f:
json.dump(res, f)
def save_act(self, imgs, save_path, ifImage):
if ifImage:
for i, img in enumerate(imgs):
img = img.transpose(1, 2, 0)
img = img[:, :, ::-1]
cv2.imwrite(os.path.join(save_path, 'im' + str(i) + '.jpg'),
img)
else:
for i, img in enumerate(imgs):
for j, act in enumerate(img):
cv2.imwrite(
os.path.join(save_path,
'im' + str(i) + 'act' + str(j) + '.jpg'),
act.numpy() * 255)
def save_images(self, imgs, batch_index, save_path, if_original=False):
for i, img in enumerate(imgs):
img = img[:, :, ::-1] # change to BGR
#from IPython import embed
#embed()
if not if_original:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Original.jpg'),
img)
else:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Pred.jpg'), img)
def color_images(self, pred, target):
imgs = []
for p, t in zip(pred, target):
tmp = p * 2 + t
np.squeeze(tmp)
img = np.zeros((p.shape[0], p.shape[1], 3))
# bkg:negative, building:postive
#from IPython import embed
#embed()
img[np.where(tmp == 0)] = [0, 0, 0] # Black RGB, for true negative
img[np.where(tmp == 1)] = [255, 0,
0] # Red RGB, for false negative
img[np.where(tmp == 2)] = [0, 255,
0] # Green RGB, for false positive
img[np.where(tmp == 3)] = [255, 255,
0] #Yellow RGB, for true positive
imgs.append(img)
return imgs
class Trainer(object):
def __init__(self, config, args):
self.args = args
self.config = config
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
config)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': config.lr
}, {
'params': model.get_10x_lr_params(),
'lr': config.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=config.momentum,
weight_decay=config.weight_decay)
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(
weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
#cudnn.benchmark = True
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image)
loss = self.criterion(output, target)
loss.backward()
self.optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
#Acc_class = self.evaluator.Pixel_Accuracy_Class()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
#FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
print('Validation:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + image.data.shape[0]))
print("Acc:{}, IoU:{}, mIoU:{}".format(Acc, IoU, mIoU))
print('Loss: %.3f' % test_loss)
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
print('Saving state, epoch:', epoch)
torch.save(
self.model.module.state_dict(), self.args.save_folder +
'models/' + 'epoch' + str(epoch) + '.pth')
loss_file = {'Acc': Acc, 'IoU': IoU, 'mIoU': mIoU}
with open(
os.path.join(self.args.save_folder, 'eval',
'epoch' + str(epoch) + '.json'), 'w') as f:
json.dump(loss_file, f)
class Trainer(object):
def __init__(self, config, args):
self.args = args
self.config = config
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(config)
self.target_train_loader, self.target_val_loader, self.target_test_loader, _ = make_target_data_loader(config)
# Define network
self.model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
train_params = [{'params': self.model.get_1x_lr_params(), 'lr': config.lr},
{'params': self.model.get_10x_lr_params(), 'lr': config.lr * config.lr_ratio}]
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params, momentum=config.momentum,
weight_decay=config.weight_decay)
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.consistency = ConsistencyLoss(cuda=args.cuda)
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler, config.lr,
config.epochs, len(self.train_loader),
config.lr_step, config.warmup_epochs)
self.summary = TensorboardSummary('./train_log')
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.best_pred_source = 0.0
# Resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint, map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, args.start_epoch))
def training(self, epoch):
train_loss, seg_loss_sum, consistency_loss_sum = 0.0, 0.0, 0.0
self.model.train()
if config.freeze_bn:
self.model.module.freeze_bn()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
target_train_iterator = iter(self.target_train_loader)
for i, sample in enumerate(tbar):
itr = epoch * len(self.train_loader) + i
self.summary.writer.add_scalar('Train/lr', self.optimizer.param_groups[0]['lr'], itr)
A_image, A_target = sample['image'], sample['label']
# Get one batch from target domain
try:
target_sample = next(target_train_iterator)
except StopIteration:
target_train_iterator = iter(self.target_train_loader)
target_sample = next(target_train_iterator)
B_image, B_target, B_image_pair = target_sample['image'], target_sample['label'], target_sample['image_pair']
if self.args.cuda:
A_image, A_target = A_image.cuda(), A_target.cuda()
B_image, B_target, B_image_pair = B_image.cuda(), B_target.cuda(), B_image_pair.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred_source, 0., self.config.lr_ratio)
A_output, A_feat, A_low_feat = self.model(A_image)
B_output, B_feat, B_low_feat = self.model(B_image)
B_output_pair, B_feat_pair, B_low_feat_pair = self.model(B_image_pair)
self.optimizer.zero_grad()
# Train seg network
# Supervised loss
seg_loss = self.criterion(A_output, A_target)
main_loss = seg_loss
# Consistency loss
consistency_loss = 0.01 * self.ConsistencyLoss(B_output, B_output_pair)
main_loss += consistency_loss
main_loss.backward()
self.optimizer.step()
seg_loss_sum += seg_loss.item()
consistency_loss_sum += consistency_loss.item()
train_loss += seg_loss.item()
self.summary.writer.add_scalar('Train/SegLoss', seg_loss.item(), itr)
self.summary.writer.add_scalar('Train/ConsistencyLoss', consistency_loss.item(), itr)
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
# Show the results of the last iteration
#if i == len(self.train_loader)-1:
print("Add Train images at epoch"+str(epoch))
self.summary.visualize_image('Train-Source', self.config.dataset, A_image, A_target, A_output, epoch, 5)
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.config.batch_size + A_image.data.shape[0]))
print('Loss: %.3f' % train_loss)
def validation(self, epoch):
def get_metrics(tbar, if_source=False):
self.evaluator.reset()
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output, low_feat, feat = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target_ = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target_, pred)
if if_source:
print("Add Validation-Source images at epoch"+str(epoch))
self.summary.visualize_image('Val-Source', self.config.dataset, image, target, output, epoch, 5)
else:
print("Add Validation-Target images at epoch"+str(epoch))
self.summary.visualize_image('Val-Target', self.config.target, image, target, output, epoch, 5)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
if if_source:
print('Validation on source:')
else:
print('Validation on target:')
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.config.batch_size + image.data.shape[0]))
print("Acc:{}, IoU:{}, mIoU:{}".format(Acc, IoU, mIoU))
print('Loss: %.3f' % test_loss)
if if_source:
names = ['source', 'source_acc', 'source_IoU', 'source_mIoU']
self.summary.writer.add_scalar('Val/SourceAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/SourceIoU', IoU, epoch)
else:
names = ['target', 'target_acc', 'target_IoU', 'target_mIoU']
self.summary.writer.add_scalar('Val/TargetAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/TargetIoU', IoU, epoch)
return Acc, IoU, mIoU
self.model.eval()
tbar_source = tqdm(self.val_loader, desc='\r')
s_acc, s_iou, s_miou = get_metrics(tbar_source, True)
new_pred_source = s_iou
if new_pred_source > self.best_pred_source:
is_best = True
self.best_pred_source = max(new_pred_source, self.best_pred_source)
print('Saving state, epoch:', epoch)
torch.save(self.model.module.state_dict(), self.args.save_folder + 'models/'
+ 'epoch' + str(epoch) + '.pth')
loss_file = {'s_Acc': s_acc, 's_IoU': s_iou, 's_mIoU': s_miou}
with open(os.path.join(self.args.save_folder, 'eval', 'epoch' + str(epoch) + '.json'), 'w') as f:
json.dump(loss_file, f)
class Test:
def __init__(self,
model_path,
config,
bn,
save_path,
save_batch,
sample_number,
trial=100,
cuda=False,
num_layers=61,
city_name='Vegas'):
self.bn = bn
#self.target=config.all_dataset
#self.target.remove(config.dataset)
self.target = city_name
self.sample_number = sample_number
# load source domain
#self.source_set = spacenet.Spacenet(city=config.dataset, split='test', img_root=config.img_root, needs to be changed)
#self.source_loader = DataLoader(self.source_set, batch_size=16, shuffle=False, num_workers=2)
self.source_loader = None
self.save_path = save_path
self.save_batch = save_batch
self.trial = trial
self.target_set = []
self.target_loader = []
self.target_trainset = []
self.target_trainloader = []
self.config = config
self.num_layers = num_layers
# load other domains
city = self.target
test = spacenet.Spacenet(city=city,
split='val',
img_root=config.img_root,
gt_root=config.gt_root,
mean_std=config.mean_std,
if_augment=config.if_augment,
repeat_count=config.repeat_count)
self.target_set.append(test)
self.target_loader.append(
DataLoader(test, batch_size=16, shuffle=False, num_workers=2))
train = spacenet.Spacenet(city=city,
split='train',
img_root=config.img_root,
gt_root=config.gt_root,
mean_std=config.mean_std,
if_augment=config.if_augment,
repeat_count=config.repeat_count,
sample_number=sample_number)
self.target_trainset.append(train)
self.target_trainloader.append(
DataLoader(train, batch_size=16, shuffle=False, num_workers=2))
self.model = DeepLab(num_classes=2,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
if cuda:
self.checkpoint = torch.load(model_path)
else:
self.checkpoint = torch.load(model_path,
map_location=torch.device('cpu'))
#print(self.checkpoint.keys())
self.model.load_state_dict(self.checkpoint['model'])
self.evaluator = Evaluator(2)
self.cuda = cuda
if cuda:
self.model = self.model.cuda()
def get_source_bn(self, trainloader, model, city):
dirname = os.path.join(self.save_path, city + '_bn')
bn_before = {}
model.eval()
for sample in trainloader:
image, target, path = sample['image'], sample['label'], sample[
'path']
if self.cuda:
image, target = image.cuda(), target.cuda()
output = model(image)
layr = 0
for h in model.modules():
if isinstance(h, nn.Conv2d):
bn_before[(layr), 'weight'] = np.squeeze(
h.weight.detach().cpu().numpy())
if isinstance(h, nn.BatchNorm2d): #Conv2d):
bn_before[(layr, 'mean')] = h.running_mean.cpu().numpy()
bn_before[(layr, 'var')] = h.running_var.cpu().numpy()
layr += 1
pickle.dump(
bn_before,
open(
os.path.join(
dirname,
'bnAll_before_{}.pickle'.format(converged_model)),
'wb'))
break
return
def get_performance(self, dataloader, trainloader, city, adabn_layer,
model):
# change mean and var of bn to adapt to the target domain
dirname = os.path.join(self.save_path, city + '_bn')
if not os.path.exists(dirname):
os.makedirs(dirname)
model.eval()
current_adabn = 0
for h in model.modules():
if isinstance(h, SynchronizedBatchNorm2d) or isinstance(
h, nn.BatchNorm2d):
current_adabn += 1
if current_adabn == adabn_layer: #current_adabn ==1 or
h.train()
if self.bn and city != self.config.dataset:
print('BN Adaptation on' + city)
for i, sample in enumerate(trainloader):
image, target, path = sample['image'], sample['label'], sample[
'path']
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = model(image)
#bn={}
#model.eval()
#for sample in trainloader:
# image, target, path = sample['image'], sample['label'], sample['path']
# if self.cuda:
# image, target = image.cuda(), target.cuda()
# output = model(image)
# layr=0
# for h in model.modules():
# if isinstance(h, nn.Conv2d):
# bn[(layr,'weight')] = np.squeeze(h.weight.detach().cpu().numpy())
# if isinstance(h, nn.BatchNorm2d): #Conv2d):
# bn[(layr,'mean')] = h.running_mean.cpu().numpy()
# bn[(layr,'var')] = h.running_var.cpu().numpy()
# layr +=1
# pickle.dump(bn,open(os.path.join(dirname, 'bnAll_after_k_1_{}_{}_{}.pickle'.format(converged_model.rstrip('.pth'), self.sample_number, adabn_layer)),'wb')); #self.trial
# break
batch = self.save_batch
model.eval()
self.evaluator.reset()
tbar = tqdm(dataloader, desc='\r')
# save in different directories
if self.bn:
save_path = os.path.join(self.save_path, city + '_bn')
else:
save_path = os.path.join(self.save_path, city)
# evaluate on the test dataset
for i, sample in enumerate(tbar):
image, target, path = sample['image'], sample['label'], sample[
'path']
#print(path)
if self.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = model(image)
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# save pictures
if batch > 0:
if not os.path.exists(self.save_path):
os.mkdir(self.save_path)
if not os.path.exists(save_path):
os.mkdir(save_path)
image = image.cpu().numpy() * 255
image = image.transpose(0, 2, 3, 1).astype(int)
imgs = self.color_images(pred, target)
self.save_images(imgs, batch, save_path, False)
self.save_images(image, batch, save_path, True)
batch -= 1
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
#state = { 'model': model.state_dict()}
#torch.save(state, os.path.join(dirname, 'weights_{}_{}_after_adabn_k_{}.pth'.format(converged_model.rstrip('.pth'), self.sample_number, adabn_layer)))
return Acc, IoU, mIoU
def test(self, name, converged_model):
main_model = self.model
source_test = 0
if source_test:
A, I, Im = self.get_performance(self.source_loader, None,
self.config.dataset)
iter_model = copy.deepcopy(main_model)
self.get_source_bn(self.target_trainloader[0], iter_model, self.target)
i_all = []
for l in range(0, self.num_layers, 1):
iter_model = copy.deepcopy(main_model)
tA, tI, tIm = self.get_performance(self.target_loader[0],
self.target_trainloader[0],
self.target, l, iter_model)
#pdb.set_trace()
i_all.append(tI)
name_f = '{}/{}_{}_adabn_k.pickle'.format(name, self.target,
converged_model)
pickle.dump((i_all, converged_model, self.sample_number),
open(name_f, 'wb'))
#pickle.dump((i_all,converged_model,self.sample_number), open(name,'wb'))
res = {}
if source_test:
print("Test for source domain:")
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(self.config.dataset, A,
I, Im))
res[config.dataset] = {'Acc': A, 'IoU': I, 'mIoU': Im}
print('Test for target domain:')
for i, city in enumerate(self.target):
print("{}: Acc:{}, IoU:{}, mIoU:{}".format(city, tA[i], tI[i],
tIm[i]))
res[city] = {'Acc': tA[i], 'IoU': tI[i], 'mIoU': tIm[i]}
if self.bn:
name = 'train_log/test_bn.json'
else:
name = 'train_log/test.json'
with open(name, 'w') as f:
json.dump(res, f)
def save_images(self, imgs, batch_index, save_path, if_original=False):
for i, img in enumerate(imgs):
#img = img[:,:,::-1] # change to BGR
#from IPython import embed
#embed()
if not if_original:
cv2.imwrite(
os.path.join(save_path,
str(batch_index) + str(i) + '_Original.jpg'),
img)
else:
img = img.astype(np.uint8)
img = Image.fromarray(img)
img.save(
os.path.join(save_path,
str(batch_index) + str(i) + '_Pred.jpg'))
#cv2.imwrite(os.path.join(save_path, str(batch_index) + str(i) + '_Pred.jpg'), img)
def color_images(self, pred, target):
imgs = []
for p, t in zip(pred, target):
tmp = p * 2 + t
np.squeeze(tmp)
img = np.zeros((p.shape[0], p.shape[1], 3))
# bkg:negative, building:postive
#from IPython import embed
#embed()
img[np.where(tmp == 0)] = [0, 0, 0] # Black RGB, for true negative
img[np.where(tmp == 1)] = [255, 0,
0] # Red RGB, for false negative
img[np.where(tmp == 2)] = [0, 255,
0] # Green RGB, for false positive
img[np.where(tmp == 3)] = [255, 255,
0] #Yellow RGB, for true positive
imgs.append(img)
return imgs
