rrtnet.py

# -*- coding: utf-8 -*-
"""RRTNet.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/18zd_wqTCw9gSCOoyrNAY3aa70B9mSmF_
"""

import numpy as np
import torch 
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
# !unzip images.zip

class RRTNet(nn.Module):
  def __init__(self, p=0.5):
    super(RRTNet, self).__init__()
    self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
    self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
    self.conv3 = nn.Conv2d(16, 32, 3, padding=1)
    self.conv4 = nn.Conv2d(32, 64, 3, padding=1)
    self.conv5 = nn.Conv2d(64, 128, 3, padding=1)
    self.fc1 = nn.Linear(4 * 4 * 128, 256)
    self.fc2 = nn.Linear(256 + 6, 2048)
    self.fc3 = nn.Linear(2048, 1024)
    self.fc4 = nn.Linear(1024, 512)
    self.fc5 = nn.Linear(512, 2)
    self.dropout = nn.Dropout(p)

  def forward(self, x):
    config_states, maps = x
    maps = maps.unsqueeze(1).float() # .to("cuda")
    config_states = config_states.float() # .to("cuda")
    x = F.max_pool2d(F.relu(self.conv1(maps)), 2)
    x = F.max_pool2d(F.relu(self.conv2(x)), 2)
    x = F.max_pool2d(F.relu(self.conv3(x)), 2)
    x = F.max_pool2d(F.relu(self.conv4(x)), 2)
    x = F.max_pool2d(F.relu(self.conv5(x)), 2)
    x = x.view(-1, self.num_flat_features(x))
    x = F.relu(self.fc1(x))
    x = torch.cat((x, config_states), axis=1)
    x = F.relu(self.fc2(x))
    x = self.dropout(x)
    x = F.relu(self.fc3(x))
    x = self.dropout(x)
    x = F.relu(self.fc4(x))
    x = self.fc5(x)
    return x
  
  def num_flat_features(self, x):
    size = x.size()[1:]  # all dimensions except the batch dimension
    num_features = 1
    for s in size:
        num_features *= s
    return num_features

# net = RRTNet(0.5).float().to("cuda")
# optimizer = optim.Adam(net.parameters(), lr=0.001)
# criterion = nn.MSELoss()
# JHist = []

# from dataloader import NonHolonomicDataset
# from torch.utils.data import DataLoader
# path = './'
# dataset = NonHolonomicDataset('data.csv', path, grayscale=True)
# dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

# train_x, train_y = None, None
# for j in range(100):
#   for i, (input_data, labels) in enumerate(dataloader):
#     y_labels = labels.float().to("cuda")
    
#     optimizer.zero_grad()
#     y_pred = net(input_data)

#     loss = criterion(y_pred, y_labels)
#     loss.backward()
    
#     optimizer.step()
#     epoch_loss = loss.item()
#     JHist.append(epoch_loss)
    
#     print(i + j * 119, epoch_loss)

# # overfit on one batch
# labels = labels.to("cuda").float()
# for i in range(5000):
#   optimizer.zero_grad()
#   y_pred = net(input_data)
#   loss = criterion(y_pred, labels)
#   loss.backward()
#   optimizer.step()
#   epoch_loss = loss.item()
#   print(i, epoch_loss)

# import matplotlib.pyplot as plt
# JHist_avg = []
# sum = 0
# for index in range(30, len(JHist)):
#   sum += JHist[index]
#   if index % 99 == 0:
#     JHist_avg.append(sum / 100)
#     sum = 0

# plt.plot(JHist_avg)

# from google.colab import files
# PATH = './RRTNet.pth'
# torch.save(net.state_dict(), PATH)
# files.download(PATH)