def midfunc():
print 'load Model....\n'
t1 = time.time()
Net = ERFNet(2)
model = './erf_save/model/main-erfnet-step-5998-epoch-20.pth'
#print list(model.children())
Net.load_state_dict(torch.load(model))
Net = Net.cuda()
print 'Done.\n'
print 'time:', time.time() - t1, '\n'
print 'compute output....\n'
infer(Net)
return 0
class Model:
def __init__(self, args):
self.args = args
self.device = args.device
if args.model == 'stackhourglass':
self.model = ERFNet(2)
self.model.cuda() #= self.model.to(self.device)
if args.use_multiple_gpu:
self.model = torch.nn.DataParallel(self.model)
self.model.cuda()
if args.mode == 'train':
self.loss_function = MonodepthLoss(
n=3,
SSIM_w=0.8,
disp_gradient_w=0.1, lr_w=1).to(self.device)
self.optimizer = optim.Adam(self.model.parameters(),
lr=args.learning_rate)
self.val_n_img, self.val_loader = prepare_dataloader(args.val_data_dir, 'test',
args.augment_parameters,
False, 1,
(args.input_height, args.input_width),
args.num_workers)
# self.model.load_state_dict(torch.load(args.loadmodel))#jiazai yuxunlian moxing jiezhu ciju
else:
self.model.load_state_dict(torch.load(args.model_path))
args.augment_parameters = None
args.do_augmentation = False
args.batch_size = 1
self.n_img, self.loader = prepare_dataloader(args.data_dir, args.mode,
args.augment_parameters,
args.do_augmentation, args.batch_size,
(args.input_height, args.input_width),
args.num_workers)
if 'cuda' in self.device:
torch.cuda.synchronize()
def disparity_loader(self,path):
return Image.open(path)
def train(self):
# losses = []
#val_losses = []
best_loss = float('Inf')
best_val_loss = 1000000.0
for epoch in range(self.args.epochs):
# losses = []
if self.args.adjust_lr:
adjust_learning_rate(self.optimizer, epoch,
self.args.learning_rate)
c_time = time.time()
running_loss = 0.0
self.model.train() #?start?
for data in self.loader: #(test-for-train)
# Load data
self.optimizer.zero_grad()
# loss=0
data = to_device(data, self.device)
left = data['left_image']
bg_image = data['bg_image']
# plt.figure()
# plt.imshow(left.squeeze().cpu().detach().numpy())
# plt.show()
# plt.imshow(bg_image.squeeze().cpu().detach().numpy())
# plt.show()
disps = self.model(left)
# print('gggggggggggggggggggggggggggggggggggggggggg',left.shape,disps.shape)
# plt.imshow(disps.squeeze().cpu().detach().numpy())
# plt.show()
loss = self.loss_function(disps,bg_image)
ssim_loss = ssim(disps,bg_image)
ssim=1-ssim_loss
loss1=loss*1+ssim*1
loss1.backward()
self.optimizer.step()
# losses.append(loss.item())
running_loss += loss1.item()
# print(' time = %.2f' %(time.time() - c_time))
running_loss /=( self.n_img / self.args.batch_size)
# print('running_loss:', running_loss)
# running_val_loss /= (self.val_n_img / self.args.batch_size)
print('Epoch:',epoch + 1,'train_loss:',running_loss,'time:',round(time.time() - c_time, 3),'s')
TrueLoss=math.sqrt( running_loss )*255
print ('TrueLoss:',TrueLoss)
if epoch%5==0:
self.model.eval()
i=0
running_val_loss = 0.0
for data in self.val_loader:
data = to_device(data, self.device)
left = data['left_image']
bg_image = data['bg_image']
with torch.no_grad():
# newinput=torch.cat([left,bg],1)
# disps = self.model(newinput)
disps = self.model(left)
loss = self.loss_function(disps,bg_image)
ssim_loss = ssim(disps,bg_image)
ssim=1-ssim_loss
loss1=loss*1+ssim*1
# loss1 = self.loss_function((disps+left1.float()),target,mask)
running_val_loss+=loss1
running_val_loss/=200
print( 'running_val_loss = %.12f' %(running_val_loss))
if running_val_loss < best_val_loss:
self.save(self.args.model_path[:-4] + '_cpt.pth')
best_val_loss = running_val_loss
print('Model_saved')
# self.save(self.args.model_path[:-4] + '_'+str(epoch+1)+'.pth')
# if epoch==100:
# self.save(self.args.model_path[:-4] + '_100.pth')
# if epoch==150:
# self.save(self.args.model_path[:-4] + '_150.pth')
# if epoch==200:
# self.save(self.args.model_path[:-4] + '_200.pth')
# if epoch==250:
# self.save(self.args.model_path[:-4] + '_250.pth')
# self.save(self.args.model_path[:-4] + '_last.pth')#上一回
# if running_loss < best_val_loss:
# #print(running_val_loss)
# #print(best_val_loss)
# self.save(self.args.model_path[:-4] + '_cpt.pth')
# best_val_loss = running_loss
# print('Model_saved')
# print ('Finished Training. Best loss:', best_loss)
#self.save(self.args.model_path)
def save(self, path):
torch.save(self.model.state_dict(), path)
def load(self, path):
self.model.load_state_dict(torch.load(path))
class Model:
def __init__(self, args):
self.args = args
self.device = args.device
if args.model == 'stackhourglass':
self.model = ERFNet(1)
self.model2 = ldyNet(1)
self.model.cuda() #= self.model.to(self.device)
self.model2.cuda()
if args.use_multiple_gpu:
self.model = torch.nn.DataParallel(self.model)
self.model.cuda()
self.model2 = torch.nn.DataParallel(self.model2)
self.model2.cuda()
if args.mode == 'train':
self.loss_function = MonodepthLoss(n=3,
SSIM_w=0.8,
disp_gradient_w=0.1,
lr_w=1).to(self.device)
self.optimizer = optim.Adam(self.model.parameters(),
lr=args.learning_rate)
self.optimizer2 = optim.Adam(self.model2.parameters(),
lr=args.learning_rate)
self.val_n_img, self.val_loader = prepare_dataloader(
args.val_data_dir, 'test', args.augment_parameters, False, 1,
(args.input_height, args.input_width), args.num_workers)
# self.model.load_state_dict(torch.load(args.loadmodel))#jiazai yuxunlian moxing jiezhu ciju
else:
self.model.load_state_dict(torch.load(args.model_path))
self.model2.load_state_dict(torch.load(args.model2_path))
args.augment_parameters = None
args.do_augmentation = False
args.batch_size = 1
self.n_img, self.loader = prepare_dataloader(
args.data_dir, args.mode, args.augment_parameters,
args.do_augmentation, args.batch_size,
(args.input_height, args.input_width), args.num_workers)
if 'cuda' in self.device:
torch.cuda.synchronize()
def disparity_loader(self, path):
return Image.open(path)
def train(self):
# losses = []
#val_losses = []
best_loss = float('Inf')
best_val_loss = 1000000.0
for epoch in range(self.args.epochs):
# losses = []
if self.args.adjust_lr:
adjust_learning_rate(self.optimizer, epoch,
self.args.learning_rate)
if self.args.adjust_lr:
adjust_learning_rate2(self.optimizer2, epoch,
self.args.learning_rate)
c_time = time.time()
running_loss = 0.0
self.model.train() #?start?
self.model2.train()
for data in self.loader: #(test-for-train)
# Load data
data = to_device(data, self.device)
left = data['left_image']
bg_image = data['bg_image']
#template=data['template']
#randomMatrix=np.random.randint(0,1,(256,256))
# Net1Iput=torch.from_numpy(randomMatrix)
disps = self.model(left)
#j=0
# while j<(self.args.batch_size-1):
# disps1=disps
i = 0
left1 = left
while i < 15:
left1 = torch.cat((left1, left), 1)
i = i + 1
Net2Iput = disps.mul(left1) #点乘
Net2Iput11, Net2Iput12 = Net2Iput.split(128, 2)
Net2Iput11, Net2Iput21 = Net2Iput11.split(128, 3)
Net2Iput12, Net2Iput22 = Net2Iput12.split(128, 3)
# Net2Iput1=Net2Iput[:127]
# Net2Iput2=Net2Iput[128:]
# Net2Iput11=Net2Iput1[:,0:127]
# Net2Iput12=Net2Iput1[:,128:255]
# Net2Iput21=Net2Iput2[:,0:127]
# Net2Iput22=Net2Iput2[:,128:255]
Net2Iput11 = torch.sum(Net2Iput11, 2, keepdim=True, out=None)
Net2Iput11 = torch.sum(Net2Iput11, 3, keepdim=True, out=None)
Net2Iput12 = torch.sum(Net2Iput12, 2, keepdim=True, out=None)
Net2Iput12 = torch.sum(Net2Iput12, 3, keepdim=True, out=None)
Net2Iput21 = torch.sum(Net2Iput21, 2, keepdim=True, out=None)
Net2Iput21 = torch.sum(Net2Iput21, 3, keepdim=True, out=None)
Net2Iput22 = torch.sum(Net2Iput22, 2, keepdim=True, out=None)
Net2Iput22 = torch.sum(Net2Iput22, 3, keepdim=True, out=None)
Net2Iput1 = torch.cat((Net2Iput11, Net2Iput12), 2)
Net2Iput2 = torch.cat((Net2Iput21, Net2Iput22), 2)
Net2Iput = torch.cat((Net2Iput1, Net2Iput2), 3)
# Net2Iput=torch.cat((Net2Iput,Net2Iput22),1)
ps = nn.PixelShuffle(4)
Net2Iput = ps(Net2Iput)
Net2Out = self.model2(Net2Iput)
self.optimizer2.zero_grad()
lossnet2 = self.loss_function(Net2Out, bg_image)
lossnet2.backward(retain_graph=True)
self.optimizer2.step()
self.optimizer.zero_grad()
# loss=0
# plt.figure()
# plt.imshow(left.squeeze().cpu().detach().numpy())
# plt.show()
# plt.imshow(bg_image.squeeze().cpu().detach().numpy())
# plt.show()
#disps = self.model(left)
# print('gggggggggggggggggggggggggggggggggggggggggg',left.shape,disps.shape)
# plt.imshow(disps.squeeze().cpu().detach().numpy())
# plt.show()
loss1 = self.loss_function(Net2Out, bg_image)
loss1.backward(retain_graph=True)
self.optimizer.step()
# losses.append(loss.item())
running_loss += lossnet2.item()
# print(' time = %.2f' %(time.time() - c_time))
running_loss /= (self.n_img / self.args.batch_size)
# print('running_loss:', running_loss)
# running_val_loss /= (self.val_n_img / self.args.batch_size)
print('Epoch:', epoch + 1, 'train_loss:', running_loss, 'time:',
round(time.time() - c_time, 3), 's')
TrueLoss = math.sqrt(running_loss) * 255
print('TrueLoss:', TrueLoss)
if epoch % 2 == 0:
#self.model.eval()
self.model.eval()
i = 0
running_val_loss = 0.0
for data in self.val_loader:
data = to_device(data, self.device)
left = data['left_image']
bg_image = data['bg_image']
#with torch.no_grad():
# newinput=torch.cat([left,bg],1)
# disps = self.model(newinput)
#disps = self.model(left)
#Net2Out=self.model2(disps)
#lossnet2 = self.loss_function(Net2Out,bg_image)
# loss1 = self.loss_function((disps+left1.float()),target,mask)
running_val_loss += lossnet2
#running_val_loss/=100
print('running_val_loss = %.12f' % (running_val_loss))
if running_val_loss < best_val_loss:
self.save(self.args.model_path[:-4] + '_cpt.pth')
self.save2(self.args.model2_path[:-4] + '_cpt.pth')
best_val_loss = running_val_loss
print('Model_saved')
# self.save(self.args.model_path[:-4] + '_'+str(epoch+1)+'.pth')
# if epoch==100:
# self.save(self.args.model_path[:-4] + '_100.pth')
# if epoch==150:
# self.save(self.args.model_path[:-4] + '_150.pth')
# if epoch==200:
# self.save(self.args.model_path[:-4] + '_200.pth')
# if epoch==250:
# self.save(self.args.model_path[:-4] + '_250.pth')
# self.save(self.args.model_path[:-4] + '_last.pth')#上一回
# if running_loss < best_val_loss:
# #print(running_val_loss)
# #print(best_val_loss)
# self.save(self.args.model_path[:-4] + '_cpt.pth')
# best_val_loss = running_loss
# print('Model_saved')
# print ('Finished Training. Best loss:', best_loss)
#self.save(self.args.model_path)
def save(self, path):
torch.save(self.model.state_dict(), path)
def save2(self, path):
torch.save(self.model2.state_dict(), path)
def load(self, path):
self.model.load_state_dict(torch.load(path))
class Model:
def __init__(self, args):
self.args = args
self.device = args.device
if args.model == 'stackhourglass':
self.model = ERFNet(2)
self.model.cuda()
self.model.load_state_dict(torch.load(args.loadmodel))
# print(self.model.state_dict())
# print('kkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkk')
self.n_img, self.loader = prepare_dataloader(args.data_dir, args.mode,
args.augment_parameters,
args.do_augmentation, args.batch_size,
(args.input_height, args.input_width),
args.num_workers)
if 'cuda' in self.device:
torch.cuda.synchronize()
def test(self):
i = 0
sum_loss = 0.0
sum_ssim = 0.0
average_ssim = 0.0
average_loss = 0.0
for epoch in range(self.args.epochs):
self.model.eval() #?start?
for data in self.loader: #(test-for-train)
i = i + 1
data = to_device(data, self.device)
left = data['left_image']
bg_image = data['bg_image']
disps = self.model(left)
print(disps.shape)
l_loss = l1loss(disps,bg_image)
ssim_loss = ssim(disps,bg_image)
psnr222 = psnr1(disps,bg_image)
sum_loss = sum_loss + l_loss.item()
sum_ssim += ssim_loss.item()
average_ssim = sum_ssim / i
average_loss = sum_loss / i
# print average_loss
disp_show = disps.squeeze()
bg_show = bg_image.squeeze()
print(bg_show.shape)
plt.figure()
plt.subplot(1,2,1)
plt.imshow(disp_show.data.cpu().numpy())
plt.subplot(1,2,2)
plt.imshow(bg_show.data.cpu().numpy())
plt.show()
print('average loss',average_loss,average_ssim)
