def main(args):
model = network.net()
model = load_model(args, model)
test_data = Dataset(args)
for i, sample in enumerate(iter(test_data)):
if i > 10:
break
postfix = sample['filename'].split('.')[-1]
if args.subnet == 'm_net':
trimap = sample['trimap']
print('sample trimap.shape: ', trimap.shape, sample['image'].shape)
trimap_softmax = torch.zeros(
[3, trimap.shape[-2], trimap.shape[-1]], dtype=torch.float32)
trimap_softmax.scatter_(0, trimap.long().data.cpu(), 1)
m_net_input = torch.cat((sample['image'], trimap_softmax), 0)
m_net_input = m_net_input.unsqueeze(0)
alpha_r = model.m_net(m_net_input)
alpha_r = alpha_r.squeeze()
print(trimap_softmax.shape, alpha_r.shape)
result = trimap_softmax[2, ...] + trimap_softmax[1, ...] * alpha_r
cv2.imwrite(os.path.join(args.save_dir, sample['filename']),
result.data.cpu().numpy() * 255)
elif args.subnet == 'end_to_end':
net_input = sample['image']
net_input = net_input.unsqueeze(0)
alpha = model(net_input)[1]
alpha = alpha.squeeze()
print('end_to_end alpha:', alpha.shape, type(alpha))
cv2.imwrite(
os.path.join(args.save_dir, sample['filename']).replace(
'.' + postfix, '_alpha.' + postfix),
alpha.data.cpu().numpy() * 255)
elif args.subnet == 't_net':
net_input = sample['image']
net_input = net_input.unsqueeze(0)
trimap = model(net_input)
trimap_softmax = F.softmax(trimap, dim=1)
print('trimap shape:', trimap_softmax.shape)
cv2.imwrite(
os.path.join(args.save_dir, sample['filename']).replace(
'.' + postfix, '_trimap.' + postfix),
np.moveaxis(trimap_softmax.squeeze().data.cpu().numpy() * 255,
(0, 1, 2), (-1, 0, 1)))
cv2.imwrite(
os.path.join(args.save_dir,
sample['filename']).replace('.' + postfix,
'_img.' + postfix),
np.moveaxis(sample['image'].data.cpu().numpy() * 255, (0, 1, 2),
(-1, 0, 1)) + (114., 121., 134.))
def main():
print("=============> Loading args")
args = get_args()
print("============> Environment init")
if args.without_gpu:
print("use CPU !")
device = torch.device('cpu')
else:
if torch.cuda.is_available():
device = torch.device('cuda')
else:
print("No GPU is is available !")
print("============> Building model ...")
model = network.net()
model.to(device)
print("============> Loading datasets ...")
train_data = getattr(dataset, args.trainData)(root_dir = args.dataDir, \
imglist = args.trainList, \
patch_size = args.patch_size)
trainloader = DataLoader(train_data,
batch_size=args.train_batch,
drop_last=True,
shuffle=True,
num_workers=args.nThreads,
pin_memory=False)
model.train()
print('============> Loss function ', args.train_phase)
print("============> Set optimizer ...")
lr = args.lr
train_params = model.parameters()
target_network = model
if args.train_phase == 'pre_train_t_net':
train_params = model.t_net.parameters()
target_network = model.t_net
elif args.train_phase == 'pre_train_m_net':
train_params = model.m_net.parameters()
target_network = model.m_net
model.t_net.eval()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, train_params), \
lr=lr, betas=(0.9, 0.999),
weight_decay=0.0005)
print("============> Start Train ! ...")
start_epoch = 1
trainlog = Train_Log(args)
if args.finetuning:
start_epoch, model = trainlog.load_model(model)
for epoch in range(start_epoch, args.nEpochs+1):
loss_ = 0
L_alpha_ = 0
L_composition_ = 0
L_cross_, L2_bg_ = 0, 0
loss_array = []
IOU_t_bg_, IOU_t_unsure_, IOU_t_fg_ = 0, 0, 0
IOU_alpha_ = 0
if args.lrdecayType != 'keep':
lr = set_lr(args, epoch, optimizer)
t0 = time.time()
for i, sample_batched in enumerate(trainloader):
print('batch ', i)
img, trimap_gt, alpha_gt = sample_batched['image'], sample_batched['trimap'], sample_batched['alpha']
img, trimap_gt, alpha_gt = img.to(device), trimap_gt.to(device), alpha_gt.to(device)
# end_to_end or pre_train_t_net
if args.train_phase != 'pre_train_m_net':
trimap_pre, alpha_pre = model(img)
loss, L_alpha, L_composition, L_cross, L2_cross, IOU_t, IOU_alpha = loss_function(args,
img,
trimap_pre,
trimap_gt,
alpha_pre,
alpha_gt)
print("Loss calculated %.4f\nL2: %.2f\nbg IOU: %.2f\nunsure IOU: %.2f\nfg IOU: %.2f"%(L_cross.item(), L2_cross.item(), IOU_t[0].item(), IOU_t[1].item(), IOU_t[2].item()))
else: # pre_train_m_net
trimap_softmax = torch.zeros([trimap_gt.shape[0], 3, trimap_gt.shape[2], trimap_gt.shape[3]], dtype=torch.float32)
trimap_softmax.scatter_(1, trimap_gt.long().data.cpu(), 1)
trimap_softmax = trimap_softmax.to(device)
#trimap_softmax = F.softmax(trimap_gt, dim=1)
bg_gt, unsure_gt, fg_gt = torch.split(trimap_softmax, 1, dim=1)
m_net_input = torch.cat((img, trimap_softmax), 1).to(device)
alpha_r = model.m_net(m_net_input)
alpha_p = fg_gt + unsure_gt * alpha_r
loss, L_alpha, L_composition, L_cross, L2_cross, IOU_t, IOU_alpha = loss_function(args,
img,
trimap_gt,
trimap_gt,
alpha_p,
alpha_gt)
print('loss: %.5f\tL_composision: %.5f\tL_alpha: %.5f'%(loss.item(), L_composition.item(), L_alpha.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_ += loss.item()
L_alpha_ += L_alpha.item()
L_composition_ += L_composition.item()
L_cross_ += L_cross.item()
L2_bg_ += L2_cross.item()
IOU_t_bg_ += IOU_t[0].item()
IOU_t_unsure_ += IOU_t[1].item()
IOU_t_fg_ += IOU_t[2].item()
IOU_alpha_ += IOU_alpha.item()
loss_array.append(loss.item())
# TENSORBOARD SCALARS
trainlog.add_scalar('loss', loss.item())
trainlog.add_scalar('T_net_loss', L_cross.item())
trainlog.add_scalar('T_net_bg_L2', L2_cross.item())
trainlog.add_scalar('M_net_alpha', L_alpha.item())
trainlog.add_scalar('M_net_composition', L_composition.item())
trainlog.add_scalar('IOU_t_bg', IOU_t[0].item())
trainlog.add_scalar('IOU_t_unsure', IOU_t[1].item())
trainlog.add_scalar('IOU_t_fg', IOU_t[2].item())
if (i+1) % 100 == 0:
for var_name, value in target_network.named_parameters():
# ignore unused parameters
if not hasattr(value.grad, 'data'):
continue
var_name = var_name.replace('.', '/')
trainlog.add_histogram(var_name, value.data.cpu().numpy())
trainlog.add_histogram(var_name+'/grad', value.grad.data.cpu().numpy())
# TENSORBOARD IMAGE
if (i+1) % 1000 == 0 and args.train_phase == 'pre_train_m_net':
trainlog.add_image('fg_gt', vutils.make_grid(fg_gt, normalize=True, nrow=4))
trainlog.add_image('unsure_gt', vutils.make_grid(unsure_gt, normalize=True, nrow=4))
trainlog.add_image('alpha_p', vutils.make_grid(alpha_p, normalize=True, nrow=4))
trainlog.add_image('alpha_r', vutils.make_grid(alpha_r, normalize=True, nrow=4))
trainlog.add_image('alpha_gt', vutils.make_grid(alpha_gt, normalize=True, nrow=4))
if (i+1) % 1000 == 0 and args.train_phase != 'pre_train_m_net':
trainlog.add_trimap(trimap_pre)
trainlog.add_trimap_gt(trimap_gt)
trainlog.add_image('origin_image', vutils.make_grid(img, normalize=True, nrow=4))
trainlog.step()
print('Done iterating all training data')
t1 = time.time()
if epoch % args.save_epoch == 0:
# speed = (t1 - t0) / 60
loss_ = loss_ / (i+1)
L_alpha_ = L_alpha_ / (i+1)
L_composition_ = L_composition_ / (i+1)
L_cross_ = L_cross_ / (i+1)
L2_bg_ = L2_bg_ / (i+1)
loss_var = np.var(loss_array)
IOU_t_bg_ = IOU_t_bg_ / (i+1)
IOU_t_unsure_ = IOU_t_unsure_ / (i+1)
IOU_t_fg_ = IOU_t_fg_ / (i+1)
IOU_alpha_ = IOU_alpha_ / (i+1)
trainlog.add_scalar('avg_loss', loss_, epoch)
trainlog.add_scalar('avg_t_loss', L_cross_, epoch)
trainlog.add_scalar('avg_t_L2_bg', L2_bg_, epoch)
trainlog.add_scalar('avg_t_loss_var', loss_var, epoch)
trainlog.add_scalar('avg_IOU_t_bg', IOU_t_bg_, epoch)
trainlog.add_scalar('avg_IOU_t_unsure', IOU_t_unsure_, epoch)
trainlog.add_scalar('avg_IOU_t_fg', IOU_t_fg_, epoch)
trainlog.add_scalar('avg_L_alpha', L_alpha_, epoch)
trainlog.add_scalar('avg_L_composition', L_composition_, epoch)
log = "[{} / {}] \tLr: {:.5f}\nloss: {:.5f}\tloss_p: {:.5f}\tloss_t: {:.5f}\tloss_var: {:.5f}\tIOU_t_bg: {:.5f}\tIOU_t_unsure: {:.5f}\tIOU_t_fg: {:.5f}\tIOU_alpha: {:.5f}\t" \
.format(epoch, args.nEpochs,
lr,
loss_,
L_alpha_+L_composition_,
L_cross_,
loss_var,
IOU_t_bg_,
IOU_t_unsure_,
IOU_t_fg_,
IOU_alpha_)
print(log)
trainlog.save_log(log)
trainlog.save_model(model, epoch)
def main():
print("=============> Loading args")
args = get_args()
print("============> Environment init")
if args.without_gpu:
print("use CPU !")
device = torch.device('cpu')
else:
if torch.cuda.is_available():
device = torch.device('cuda')
else:
print("No GPU is is available !")
print("============> Building model ...")
model = network.net()
model.to(device)
print("============> Loading datasets ...")
train_data = getattr(dataset, args.trainData)(root_dir = args.dataDir, \
imglist = args.trainList, \
patch_size = args.patch_size)
trainloader = DataLoader(train_data,
batch_size=args.train_batch,
drop_last=True,
shuffle=True,
num_workers=args.nThreads,
pin_memory=True)
print("============> Set optimizer ...")
lr = args.lr
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), \
lr=lr, betas=(0.9, 0.999),
weight_decay=0.0005)
print("============> Start Train ! ...")
start_epoch = 1
trainlog = Train_Log(args)
if args.finetuning:
start_epoch, model = trainlog.load_model(model)
model.train()
for epoch in tqdm(range(start_epoch, args.nEpochs + 1)):
loss_ = 0
L_alpha_ = 0
L_composition_ = 0
L_cross_ = 0
if args.lrdecayType != 'keep':
lr = set_lr(args, epoch, optimizer)
t0 = time.time()
for i, sample_batched in enumerate(tqdm(trainloader)):
img, trimap_gt, alpha_gt = sample_batched['image'], sample_batched[
'trimap'], sample_batched['alpha']
img, trimap_gt, alpha_gt = img.to(device), trimap_gt.to(
device), alpha_gt.to(device)
trimap_pre, alpha_pre = model(img)
loss, L_alpha, L_composition, L_cross = loss_function(
args, img, trimap_pre, trimap_gt, alpha_pre, alpha_gt)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_ += loss.item()
L_alpha_ += L_alpha.item()
L_composition_ += L_composition.item()
L_cross_ += L_cross.item()
t1 = time.time()
if epoch % args.save_epoch == 0:
# speed = (t1 - t0) / 60
loss_ = loss_ / (i + 1)
L_alpha_ = L_alpha_ / (i + 1)
L_composition_ = L_composition_ / (i + 1)
L_cross_ = L_cross_ / (i + 1)
log = "[{} / {}] \tLr: {:.5f}\nloss: {:.5f}\tloss_p: {:.5f}\tloss_t: {:.5f}\t" \
.format(epoch, args.nEpochs,
lr,
loss_,
L_alpha_+L_composition_,
L_cross_)
print(log)
trainlog.save_log(log)
trainlog.save_model(model, epoch)
def main():
print("=============> Loading args")
args = get_args()
print("============> Environment init")
if args.without_gpu:
print("use CPU !")
device = torch.device('cpu')
else:
if torch.cuda.is_available():
device = torch.device('cuda')
else:
print("No GPU is is available !")
device = torch.device('cpu')
print("============> Building model ...")
model = network.net()
model.to(device)
print("============> Loading datasets ...")
_data = getattr(dataset, args.trainData)(root_dir=args.dataDir,
imglist=args.trainList,
patch_size=args.patch_size)
lengths = [25365, 6342]
#print("train data : {}, valid data: {}".format(lengths[0],lengths[1]))
train_data, valid_data = torch.utils.data.dataset.random_split(
_data, lengths)
trainloader = DataLoader(train_data,
batch_size=args.train_batch,
drop_last=True,
shuffle=True,
num_workers=args.nThreads,
pin_memory=True)
validloader = DataLoader(valid_data,
batch_size=args.train_batch,
drop_last=True,
shuffle=True,
num_workers=args.nThreads,
pin_memory=True)
print("============> Set optimizer ...")
lr = args.lr
optimizer = optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=lr,
betas=(0.9, 0.999),
weight_decay=0.0005)
print("============> Start Train ! ...")
start_epoch = 1
trainlog = Train_Log(args)
if args.finetuning:
start_epoch, model = trainlog.load_model(model)
start_epoch += 1
model.train()
for epoch in range(start_epoch, args.nEpochs + 1):
loss_ = 0
L_alpha_ = 0
L_composition_ = 0
L_cross_ = 0
if args.lrdecayType != 'keep':
lr = set_lr(args, epoch, optimizer)
t0 = time.time()
model.train()
for i, sample_batched in enumerate(trainloader):
img, trimap_gt, alpha_gt = sample_batched['image'], sample_batched[
'trimap'], sample_batched[
'alpha'] #[8, 3, 256, 256], [8, 1, 256, 256], [8, 1, 256, 256]
img, trimap_gt, alpha_gt = img.to(device), trimap_gt.to(
device), alpha_gt.to(device)
trimap_pre, alpha_pre = model(
img) #[8, 3, 256, 256], [8, 1, 256, 256]
loss, L_alpha, L_composition, L_cross = loss_function(
args, img, trimap_pre, trimap_gt, alpha_pre, alpha_gt)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_ += loss.item()
L_alpha_ += L_alpha.item()
L_composition_ += L_composition.item()
L_cross_ += L_cross.item()
t1 = time.time()
# test
model.eval()
t_loss_ = 0
t_L_alpha_ = 0
t_L_composition_ = 0
t_L_cross_ = 0
with torch.no_grad():
for sample_batched in validloader:
img, trimap_gt, alpha_gt = sample_batched[
'image'], sample_batched['trimap'], sample_batched['alpha']
img, trimap_gt, alpha_gt = img.to(device), trimap_gt.to(
device), alpha_gt.to(device)
trimap_pre, alpha_pre = model(img)
loss, L_alpha, L_composition, L_cross = loss_function(
args, img, trimap_pre, trimap_gt, alpha_pre, alpha_gt)
t_loss_ += loss.item()
t_L_alpha_ += L_alpha.item()
t_L_composition_ += L_composition.item()
t_L_cross_ += L_cross.item()
if epoch % args.save_epoch == 0:
# speed = (t1 - t0) / 60
train_len = lengths[0]
test_len = lengths[1]
loss_ = loss_ / train_len
L_alpha_ = L_alpha_ / train_len
L_composition_ = L_composition_ / train_len
L_cross_ = L_cross_ / train_len
t_loss_ = t_loss_ / test_len
t_L_alpha_ = t_L_alpha_ / test_len
t_L_composition_ = t_L_composition_ / test_len
t_L_cross_ = t_L_cross_ / test_len
log = "[{} / {}] \tLr: {:.5f}\nloss: {:.5f}\tloss_p: {:.5f}\tloss_t: {:.5f}\t" \
.format(epoch, args.nEpochs, lr, loss_, L_alpha_+L_composition_, L_cross_)
print("KKTKKTKKTKKT save_log 전 epoch: " + str(epoch))
t_log = "[Validation] \tloss: {:.5f}\tloss_p: {:.5f}\tloss_t: {:.5f}\t" \
.format( t_loss_, t_L_alpha_+t_L_composition_, t_L_cross_)
print(log)
print(t_log)
trainlog.save_log(log)
trainlog.save_model(model, epoch)
def main():
print("=============> Loading args")
args = get_args()
print("============> Environment init")
if args.without_gpu:
print("use CPU !")
device = torch.device('cpu')
else:
if torch.cuda.is_available():
device = torch.device('cuda')
else:
print("No GPU is is available !")
print("============> Building model ...")
model = network.net()
model.to(device)
print("============> Loading datasets ...")
_data = getattr(dataset, args.trainData)(root_dir = args.dataDir, \
imglist = args.trainList, \
patch_size = args.patch_size)
train_size = int(0.8 * len(_data))
test_size = len(_data) - train_size
lengths = [train_size, test_size]
#print("train data : {}, valid data: {}".format(lengths[0],lengths[1]))
train_data, valid_data = torch.utils.data.dataset.random_split(
_data, lengths)
print(len(train_data), len(valid_data.indices))
trainloader = DataLoader(train_data,
batch_size=args.train_batch,
drop_last=True,
shuffle=True,
num_workers=args.nThreads,
pin_memory=True)
validloader = DataLoader(valid_data,
batch_size=args.train_batch,
drop_last=True,
shuffle=True,
num_workers=args.nThreads,
pin_memory=True)
print("============> Set optimizer ...")
lr = args.lr
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), \
lr=lr, betas=(0.9, 0.999),
weight_decay=0.0005)
print("============> Start Train ! ...")
start_epoch = 1
trainlog = Train_Log(args)
if args.finetuning:
start_epoch, model = trainlog.load_model(model)
model.train()
# make lists that contain loss value of training set and validation set!!
train_loss_per_epoch = list()
val_loss_per_epoch = list()
for epoch in range(start_epoch, args.nEpochs + 1):
loss_ = 0
L_alpha_ = 0
L_composition_ = 0
L_cross_ = 0
if args.lrdecayType != 'keep':
lr = set_lr(args, epoch, optimizer)
t0 = time.time()
model.train()
for i, sample_batched in enumerate(trainloader):
img, trimap_gt, alpha_gt = sample_batched['image'], sample_batched[
'trimap'], sample_batched['alpha']
img, trimap_gt, alpha_gt = img.to(device), trimap_gt.to(
device), alpha_gt.to(device)
trimap_pre, alpha_pre = model(img)
loss, L_alpha, L_composition, L_cross = loss_function(
args, img, trimap_pre, trimap_gt, alpha_pre, alpha_gt)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_ += loss.item()
L_alpha_ += L_alpha.item()
L_composition_ += L_composition.item()
L_cross_ += L_cross.item()
t1 = time.time()
# test
model.eval()
t_loss_ = 0
t_L_alpha_ = 0
t_L_composition_ = 0
t_L_cross_ = 0
with torch.no_grad():
for sample_batched in validloader:
img, trimap_gt, alpha_gt = sample_batched[
'image'], sample_batched['trimap'], sample_batched['alpha']
img, trimap_gt, alpha_gt = img.to(device), trimap_gt.to(
device), alpha_gt.to(device)
trimap_pre, alpha_pre = model(img)
loss, L_alpha, L_composition, L_cross = loss_function(
args, img, trimap_pre, trimap_gt, alpha_pre, alpha_gt)
t_loss_ += loss.item()
t_L_alpha_ += L_alpha.item()
t_L_composition_ += L_composition.item()
t_L_cross_ += L_cross.item()
if epoch % args.save_epoch == 0:
# speed = (t1 - t0) / 60
train_len = lengths[0]
test_len = lengths[1]
loss_ = loss_ / train_len
L_alpha_ = L_alpha_ / train_len
L_composition_ = L_composition_ / train_len
L_cross_ = L_cross_ / train_len
t_loss_ = t_loss_ / test_len
t_L_alpha_ = t_L_alpha_ / test_len
t_L_composition_ = t_L_composition_ / test_len
t_L_cross_ = t_L_cross_ / test_len
log = "[{} / {}] \tLr: {:.5f}\nloss: {:.5f}\tloss_p: {:.5f}\tloss_t: {:.5f}\t" \
.format(epoch, args.nEpochs,
lr,
loss_,
L_alpha_+L_composition_,
L_cross_)
t_log = "[Validation] \tloss: {:.5f}\tloss_p: {:.5f}\tloss_t: {:.5f}\t" \
.format( t_loss_,
t_L_alpha_+t_L_composition_,
t_L_cross_)
print(log)
print(t_log)
trainlog.save_log(log)
trainlog.save_model(model, epoch)
train_loss_per_epoch.append(
loss_) # save train loss value into the list
val_loss_per_epoch.append(
t_loss_) # save validation loss value into the list
# draw the loss graph per epoch
plt.figure(figsize=(10, 10), dpi=160)
plt.plot()
plt.suptitle('loss of training and validation set')
plt.plot(list(range(1, args.nEpochs + 1)),
train_loss_per_epoch,
color='red',
linestyle='-',
label='train')
plt.plot(list(range(1, args.nEpochs + 1)),
val_loss_per_epoch,
color='red',
linestyle='--',
label='validation')
plt.text(x=0.5, y=0.5, s='learning rate = 0.0003\nEpochs = 100')
plt.xlabel("epochs")
plt.ylabel("loss")
plt.legend(loc='upper right')