graph-encoder-decoder_L6_G3.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import glob as glob
from torch.autograd import Variable
from reader import get_depth_data
from PIL import Image

path = './checkpoints/GraphEncoderDecoder/L6_G3/'

Train = np.load('iccv_dataset_train.npy')
Val = np.load('iccv_dataset_val.npy')
Test = np.load('iccv_dataset_test.npy')

class GraphEncoderDecoder(nn.Module):

    def __init__(self):
        super(GraphEncoderDecoder, self).__init__()
        self.cnn1 = nn.Conv2d(3,10,kernel_size=(8,8))
        self.pool1 = nn.MaxPool2d(2,return_indices=True)
        self.cnn2 = nn.Conv2d(10,25,kernel_size=(8,8))
        self.pool2 = nn.MaxPool2d(2,return_indices=True)
        self.cnn3 = nn.Conv2d(25,30,kernel_size=(5,5))
        self.pool3 = nn.MaxPool2d(2,return_indices=True)
        self.cnn4 = nn.Conv2d(30,30,kernel_size=(5,5))
        self.pool4 = nn.MaxPool2d(2,return_indices=True)
        self.cnn5 = nn.Conv2d(30,40,kernel_size=(3,3))
        self.pool5 = nn.MaxPool2d(2,return_indices=True)
        self.cnn6 = nn.Conv2d(40,50,kernel_size=(3,3))
        self.pool6 = nn.MaxPool2d(2,return_indices=True)

        self.unpool1 = nn.MaxUnpool2d(2)
        self.decnn1 = nn.ConvTranspose2d(50,40,kernel_size=(3,3))
        self.unpool2 = nn.MaxUnpool2d(2)
        self.decnn2 = nn.ConvTranspose2d(40,30,kernel_size=(3,3))
        self.unpool3 = nn.MaxUnpool2d(2)
        self.decnn3 = nn.ConvTranspose2d(30,30,kernel_size=(5,5))
        self.unpool4 = nn.MaxUnpool2d(2)
        self.decnn4 = nn.ConvTranspose2d(30,25,kernel_size=(5,5))
        self.unpool5 = nn.MaxUnpool2d(2)
        self.decnn5 = nn.ConvTranspose2d(25,10,kernel_size=(8,8))
        self.unpool6 = nn.MaxUnpool2d(2)
        self.decnn6 = nn.ConvTranspose2d(10,1,kernel_size=(8,8))

        self.a = nn.Linear(2*12,1)
        self.W = nn.Linear(12,12)

        self.a2 = nn.Linear(2*12,1)
        self.W2 = nn.Linear(12,12)

        self.a3 = nn.Linear(2*12,1)
        self.W3 = nn.Linear(12,12)

        self.P = nn.Linear(12,100)
        self.hidden = nn.Linear(1500,1400)

    def forward(self,x,g):
        g = [torch.Tensor([gi]).cuda() for gi in g]
        x = F.relu(self.cnn1(x))
        s1 = x.size()
        x, i1 = self.pool1(x)
        x = F.relu(self.cnn2(x))
        s2 = x.size()
        x, i2 = self.pool2(x)
        x = F.relu(self.cnn3(x))
        s3 = x.size()
        x, i3 = self.pool3(x)
        x = F.relu(self.cnn4(x))
        s4 = x.size()
        x, i4 = self.pool4(x)
        x = F.relu(self.cnn5(x))
        s5 = x.size()
        x, i5 = self.pool5(x)
        x = F.relu(self.cnn6(x))
        s6 = x.size()
        x, i6 = self.pool6(x)
        presize = x.size()
        
        h = []
        for i,hi in enumerate(g):
            a_weights = []
            h_neighbors = []
            for j, hj in enumerate(g):
                a = F.relu(self.a(torch.cat((self.W(hi),self.W(hj)),dim=1)))
                a_weights.append(a)
                h_neighbors.append(hj)
            s = []
            for aw, hn in zip(a_weights,h_neighbors):
                s.append(self.W(aw*hn))
            h.append(sum(s))
        h = torch.cat(h,dim=0)
        
        h = list(h.unsqueeze(dim=1))
        h2 = []
        for i,hi in enumerate(h):
            a_weights = []
            h_neighbors = []
            for j, hj in enumerate(h):
                a = F.relu(self.a2(torch.cat((self.W2(hi),self.W2(hj)),dim=1)))
                a_weights.append(a)
                h_neighbors.append(hj)
            s = []
            for aw, hn in zip(a_weights,h_neighbors):
                s.append(self.W2(aw*hn))
            h2.append(sum(s))
        h2 = torch.cat(h2,dim=0)
        
        h2 = list(h2.unsqueeze(dim=1))
        h3 = []
        for i,hi in enumerate(h2):
            a_weights = []
            h_neighbors = []
            for j, hj in enumerate(h2):
                a = F.relu(self.a3(torch.cat((self.W3(hi),self.W3(hj)),dim=1)))
                a_weights.append(a)
                h_neighbors.append(hj)
            s = []
            for aw, hn in zip(a_weights,h_neighbors):
                s.append(self.W3(aw*hn))
            h3.append(sum(s))
        h3 = torch.cat(h3,dim=0)

        latent = self.P(h3)
        latent = torch.mean(latent,dim=0)
        x = torch.cat((x.flatten(),latent),dim=0)
        x = self.hidden(x)
        x = x.view(presize)

        x = self.unpool1(x,i6,s6)
        x = F.relu(self.decnn1(x))
        x = self.unpool2(x,i5,s5)
        x = F.relu(self.decnn2(x))
        x = self.unpool3(x,i4,s4)
        x = F.relu(self.decnn3(x))
        x = self.unpool4(x,i3,s3)
        x = F.relu(self.decnn4(x))
        x = self.unpool5(x,i2,s2)
        x = F.relu(self.decnn5(x))
        x = self.unpool6(x,i1,s1)
        x = F.relu(self.decnn6(x))
        return x

model = GraphEncoderDecoder().cuda()
learning_rate = 1e-3
loss_fn = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
max_epochs = 15

import time
candidate_models = []
validation_losses = []
# Training and Validation
for epoch in range(max_epochs):
    epoch_loss = []
    start = time.time()
    for train in Train:
        dirs = glob.glob(train)
        gdata,data = get_depth_data(dirs)
        gd = [g[0] for g in gdata]
        x,y = np.array(Image.open(data[0][0])),data[0][1]
        x,y = x.reshape(1,3,480,640),y.reshape(1,1,480,640)
        x = Variable(torch.Tensor(x).cuda(), requires_grad=True)
        y = Variable(torch.Tensor(y).cuda(), requires_grad=False)
        pred = model(x,gd)
        loss = loss_fn(pred,y)
        epoch_loss.append(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Validation
    validation_loss = []
    for val in Val:
        val_x,val_y = [],[]
        dirs = glob.glob(val)
        g_data, data = get_depth_data(dirs)
        gd = [g[0] for g in g_data]
        val_x, val_y = np.array(Image.open(data[0][0])), data[0][1]
        val_x, val_y = val_x.reshape(1,3,480,640), val_y.reshape(1,1,480,640)
        val_x = Variable(torch.Tensor(val_x).cuda(),requires_grad=False)
        val_pred = model.forward(val_x,gd)
        val_y = Variable(torch.Tensor(val_y).cuda(),requires_grad=False)
        val_loss = loss_fn(val_pred,val_y)
        validation_loss.append(val_loss.item())

    validation_losses.append(np.array(validation_loss).mean())

    end = time.time()
    model_path = path + 'GraphEncoderDecoder-G3-6_layer-epoch_'+str(epoch)+'.model'
    torch.save(model.state_dict(), model_path) 
    candidate_models.append(model_path) 
    print('epoch loss: ' + str(np.array(epoch_loss).mean()) + ', Val loss: ' + str(np.array(validation_loss).mean()) + ', Time: ' + str((end-start)))

    if len(validation_losses) > 1:
        check = (((validation_losses[-2] - validation_losses[-1])/(validation_losses[-2])) * 100)
        if check < 1.0 and check > 0:
            break
        if check < 0:
            candidate_models.pop()
            break
# Test
test_model = GraphEncoderDecoder().cuda()
test_model.load_state_dict(torch.load(candidate_models[-1]))
test_loss = []
for test in Test:
    test_x,test_y = [],[]
    dirs = glob.glob(test)
    g_data, data = get_depth_data(dirs)
    gd = [g[0] for g in g_data]
    test_x, test_y = np.array(Image.open(data[0][0])), data[0][1]
    test_x, test_y = test_x.reshape(1,3,480,640), test_y.reshape(1,1,480,640)
    test_x = Variable(torch.Tensor(test_x).cuda(),requires_grad=False)
    test_pred = test_model.forward(test_x,gd).detach().cpu().numpy()
    test_pred = test_pred.reshape(480,640)
    for box in g_data:
        info,dis = box
        info = [int(i) for i in info[:4]]
        # str(startX),str(startY),str(endX),str(endY)
        startX ,startY, endX, endY = info[0], info[1], info[2], info[3]
        pred_dis = test_pred[startY:endY,startX:endX]/1000.0
        pred_dis = pred_dis[pred_dis > 0]
        dis_error = abs(pred_dis.mean() - dis)
        test_loss.append(dis_error)
np.save(path+'test_loss.npy',test_loss)
print('Test Loss:',np.mean(test_loss))