Commit 5d40a80d authored by Yuxuan Yang's avatar Yuxuan Yang
Browse files

train and generate rollouts via IN-Bilstm pt for relative and full pairwise relation state

parent f60fa785
import numpy as np
import torch
from torch.utils.data import Dataset
class RopeDataset(Dataset):
def __init__(self, data):
self.len = data.shape[0]*(data.shape[1]-1)
particle_num = data.shape[-2]
self.x_data = data[:,:-1,:,:].reshape([self.len,particle_num,-1]).astype(np.float32)
self.y_data = data[:,1:,:,:].reshape([self.len,particle_num,-1]).astype(np.float32)
# print("init dataset, shape: ", data.shape, "x reshape: ", self.x_data.shape,"self.len",self.len)
def __getitem__(self,index):
return self.x_data[index], self.y_data[index]
def __len__(self):
return self.len
\ No newline at end of file
import os
import random
import time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
# import kornia
from utils import *
from torch.utils.tensorboard import SummaryWriter
import tqdm
from RopeDataset import RopeDataset
from models import InteractionNetwork_bilstm, InteractionNetwork_bilstm_full
from tqdm import tqdm
import argparse
def gen_data(model, criterion, simulate_time_step, data_gt, data_gt_idx, stat_r, stat_v):
data_gt_denorm = data_gt
batch_data_pred = data_gt_denorm[:,0,:,:]
data_pred_denorm = np.zeros((data_gt_denorm.shape[0],simulate_time_step+1,data_gt_denorm.shape[2],data_gt_denorm.shape[3]))
data_pred_denorm[:,0,:,:] = data_gt_denorm[:,0,:,:]
losses = 0
for i in range(0,simulate_time_step):
label_data = data_gt[:,i+1,:,:]
label_data = torch.Tensor(label_data).cuda()
batch_data_pred = torch.Tensor(batch_data_pred).cuda()
object_for_relation_net, object_for_object_net, sender_relations, receiver_relations, target = prepare_data_pt(1, batch_data_pred, label_data, n_objects, n_relations)
predicted = model(object_for_relation_net,object_for_object_net, sender_relations, receiver_relations,stat_r, stat_v, 0)
target = (target - stat_v[0,4:])/stat_v[1,4:]
# denormalize
predicted_denorm = predicted * stat_v[1,4:] + stat_v[0,4:]
predicted_denorm = np.array(predicted_denorm.cpu().detach().numpy())
loss = criterion(predicted, target)
losses+= loss.item()
# update
data_pred_denorm[0,i+1,:,:4] = batch_data_pred[0,:,:4].detach().cpu().numpy() + predicted_denorm
data_pred_denorm[0,i+1,:,4:] = predicted_denorm
# update constrained nodes
data_pred_denorm[0,i+1,-2:] = data_gt[:,i+1,-2:]
data_pred_denorm[0,i+1,:2] = data_gt[:,i+1,:2]
batch_data_pred = data_pred_denorm[:,i+1]
print("aver loss:",losses/simulate_time_step)
return data_pred_denorm,losses/simulate_time_step
if __name__ == '__main__':
# --- load data ---
parser = argparse.ArgumentParser()
parser.add_argument('--pair_wise_type', default='rel', help='relative state| full state concatenate, pair-wise relation type') # 'rel' or 'full'
parser.add_argument('--num_rnnlayer', type=int, default=2, help="the number of bilstm layers")
args = parser.parse_args()
timestep_per_scenario = 300
timestep_per_scenario_train = timestep_per_scenario-1
train_scenario = 2700
DATA_PATH = os.getcwd()+'/data/'
data_file = DATA_PATH + 'data_pt.npy'
data = np.load(data_file)
train_data = data[0:train_scenario]
valid_data = data[train_scenario:]
print('training_data', train_data.shape)
print('valid_data',valid_data.shape)
stat_r = np.load(DATA_PATH+"relation_vector_stat_pt.npy") # stat of relative position, relative twist, relative delta pos, relative delta twist
stat_v = np.load(DATA_PATH+"v_stat_pt.npy") # stat of position, twist, delta pos, delta twist
#################
simulate_time_step = 299
num_layer = args.num_rnnlayer # 1,2,3
hidden_size = 150
n_objects = 16+1
n_relations = (n_objects-1)*2
if args.pair_wise_type == 'rel':
model_parameter_path = os.getcwd()+'/trained_models/inbilstm_pt.pth'
SAVE_DATA_FOLDER = os.getcwd()+"/generated_data/inbilstm_pt"
object_net_dim = 4+2 # [dx,dy,dz, dtwist, attri(2)]
relation_net_dim = 8+2*2 # [x,y,z,twist, dx,dy,dz,dtwist, attri(2)]
network_chosen = InteractionNetwork_bilstm
elif args.pair_wise_type == 'full':
model_parameter_path = os.getcwd()+'/trained_models/inbilstm_pt_full.pth'
SAVE_DATA_FOLDER = os.getcwd()+"/generated_data/inbilstm_pt_full"
object_net_dim = 4+2 # [dx,dy,dz, dtwist, attri(2)]
relation_net_dim = 8*2+2*2 # [x,y,z,twist, dx,dy,dz,dtwist, attri(2)]*2
network_chosen = InteractionNetwork_bilstm_full
stat_r = stat_v
else:
raise AssertionError("Unsupported pair-wise relation type, either relative ('rel') or full ('full')")
network = network_chosen(relation_net_dim, object_net_dim, hidden_size, num_layer)
network.load_state_dict(torch.load(model_parameter_path))
network.eval()
network = network.cuda()
# criterion = nn.L1Loss()
criterion = nn.MSELoss()
criterion = criterion.cuda()
stat_r_ts = torch.Tensor(stat_r).cuda()
stat_v_ts = torch.Tensor(stat_v).cuda()
os.system('mkdir -p ' + SAVE_DATA_FOLDER)
data_generated_all = np.zeros([300,simulate_time_step+1,n_objects,8])
losses = []
for i, scenrio in tqdm(enumerate(list(range(300)))):
print('generating secnario',scenrio)
start_step = 0 #
data_gt = np.array([valid_data[scenrio,start_step:start_step+simulate_time_step+1,:,:]])
data_gt_idx = [scenrio]
data_generated, loss = gen_data(network, criterion, simulate_time_step, data_gt, data_gt_idx, stat_r_ts, stat_v_ts)
data_generated_all[i,:,:,:]=data_generated[0].copy()
losses.append(loss)
np.save(SAVE_DATA_FOLDER+"data_generated_all_valid_rollout",data_generated_all)
\ No newline at end of file
import numpy as np
import torch
import torch.nn as nn
class RelationalModel(nn.Module):
def __init__(self, input_size, output_size, hidden_size):
super(RelationalModel, self).__init__()
self.output_size = output_size
self.layers = nn.Sequential(
nn.Linear(input_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, output_size),
nn.ReLU()
)
def forward(self, x):
'''
Args:
x: [batch_size, n_relations, input_size]
Returns:
[batch_size, n_relations, output_size]
'''
batch_size, n_relations, input_size = x.size()
x = x.view(-1, input_size)
x = self.layers(x)
x = x.view(batch_size, n_relations, self.output_size)
return x
class ParticleEncode(nn.Module):
def __init__(self, input_size, output_size, hidden_size):
super(ParticleEncode, self).__init__()
self.output_size = output_size
self.layers = nn.Sequential(
nn.Linear(input_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, output_size),
nn.ReLU()
)
def forward(self, x):
'''
Args:
x: [batch_size, n_relations, input_size]
Returns:
[batch_size, n_relations, output_size]
'''
batch_size, n_relations, input_size = x.size()
x = x.view(-1, input_size)
x = self.layers(x)
x = x.view(batch_size, n_relations, self.output_size)
return x
class ParticlePredictor(nn.Module):
def __init__(self, input_size, hidden_size, output_size, residual=False):
super(ParticlePredictor, self).__init__()
self.residual = residual
self.input_size = input_size
self.hidden_size = hidden_size
self.output_size = output_size
self.linear_0 = nn.Linear(input_size, hidden_size)
self.linear_1 = nn.Linear(hidden_size, hidden_size)
self.linear_2 = nn.Linear(hidden_size, hidden_size)
self.linear_3 = nn.Linear(hidden_size, output_size)
self.relu = nn.ReLU()
def forward(self, x, res=None):
'''
Args:
x: [batch_size, n_particles, input_size]
Returns:
[batch_size, n_particles, output_size]
'''
B, N, D = x.size()
x = x.view(B * N, D)
x = self.linear_1(self.relu(self.linear_0(x)))
x = self.linear_2(self.relu(self.linear_1(x)))
x = self.linear_3(self.relu(x))
if self.residual:
x = x + res.view(B * N, self.output_size)
return x
class InteractionNetwork_bilstm(nn.Module):
'''
'''
def __init__(self, relation_net_dim, object_net_dim, hidden_size, num_layers):
super(InteractionNetwork_bilstm, self).__init__()
self.state_encoder_model = ParticleEncode(object_net_dim, hidden_size, hidden_size)
self.relational_model = RelationalModel(relation_net_dim, hidden_size, hidden_size)
self.particle_predictor = ParticlePredictor(
hidden_size+hidden_size*2, hidden_size, 4)
self.bi_gru = nn.LSTM(input_size=hidden_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=False,
bidirectional=True)
def forward(self, object_for_relation_net, object_for_object_net, sender_relations, receiver_relations, stat_r, stat_v, pstep, debug=False):
# print(sender_relations.shape,object_for_relation_net.shape)
ori_size = object_for_relation_net.size()
batch_size, n_objects, n_state_all = ori_size # state_all = state + state_attr
sender_relations_t = sender_relations.permute(0, 2, 1)
receiver_relations_t = receiver_relations.permute(0, 2, 1)
senders = sender_relations_t.bmm(object_for_relation_net)
receivers = receiver_relations_t.bmm(object_for_relation_net)
# relation-centric network
delta = senders[:,:,:-2] - receivers[:,:,:-2]
delta = (delta-stat_r[0])/stat_r[1]
input_for_relation_network = torch.cat([delta,senders[:,:,-2:],receivers[:,:,-2:]],2)
relation_encode = self.relational_model(input_for_relation_network)
# state-encoder network
object_for_object_net_norm = (object_for_relation_net[:,:,4:-2]- stat_v[0,4:])/stat_v[1,4:]
object_for_object_net_norm = torch.cat([object_for_object_net_norm,object_for_relation_net[:,:,-2:]],axis=-1)
encodered_state = self.state_encoder_model(object_for_object_net_norm)
# object-centric network
### propogate through bi-gru
effect_agg = receiver_relations.bmm(relation_encode)
rnn_input = effect_agg.transpose(0,1)
rnn_out, _ = self.bi_gru(rnn_input)
rnn_out = rnn_out.transpose(0,1).view(batch_size,n_objects,-1)
# print("effect_agg.size()",effect_agg.size())
predicted= self.particle_predictor(torch.cat([encodered_state, rnn_out], 2))
return predicted
class InteractionNetwork_bilstm_full(nn.Module):
'''
normalized relative state
prediction network only take velocity
only take pair-wise relation for comparision
'''
def __init__(self, relation_net_dim, object_net_dim, hidden_size, num_layers):
super(InteractionNetwork_bilstm_full, self).__init__()
self.state_encoder_model = ParticleEncode(object_net_dim, hidden_size, hidden_size)
self.relational_model = RelationalModel(relation_net_dim, hidden_size, hidden_size)
# input: (1) particle effect
self.particle_predictor = ParticlePredictor(
hidden_size+hidden_size*2, hidden_size, 4)
# self.particle_predictor = ParticlePredictor(
# 150, 150, 4, True)
self.bi_gru = nn.LSTM(input_size=hidden_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=False,
bidirectional=True)
def forward(self, object_for_relation_net, object_for_object_net, sender_relations, receiver_relations, stat_r, stat_v, pstep):
ori_size = object_for_relation_net.size()
batch_size, n_objects, n_state_all = ori_size # state_all = state + state_attr
sender_relations_t = sender_relations.permute(0, 2, 1)
receiver_relations_t = receiver_relations.permute(0, 2, 1)
n_object_for_relation_net = (object_for_relation_net[:,:,:-2] - stat_v[0]) / stat_v[1]
n_object_for_relation_net = torch.cat([n_object_for_relation_net,object_for_relation_net[:,:,-2:]],-1)
senders = sender_relations_t.bmm(n_object_for_relation_net)
receivers = receiver_relations_t.bmm(n_object_for_relation_net)
# relation-centric network
input_for_relation_network = torch.cat([senders,receivers],2)
relation_encode = self.relational_model(input_for_relation_network)
object_for_object_net_norm = (object_for_relation_net[:,:,4:-2]- stat_v[0,4:])/stat_v[1,4:]
object_for_object_net_norm = torch.cat([object_for_object_net_norm,object_for_relation_net[:,:,-2:]],axis=-1)
encodered_state = self.state_encoder_model(object_for_object_net_norm)
### propogate through bi-gru
effect_agg = receiver_relations.bmm(relation_encode)
bi_gru_input = effect_agg.transpose(0,1)
bi_out, _ = self.bi_gru(bi_gru_input)
bi_out = bi_out.transpose(0,1).view(batch_size,n_objects,-1)
# # object-centric network
predicted= self.particle_predictor(torch.cat([encodered_state, bi_out], 2))
return predicted
\ No newline at end of file
import os
import random
import time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
# import kornia
from utils import *
from torch.utils.tensorboard import SummaryWriter
import tqdm
from RopeDataset import RopeDataset
from models import InteractionNetwork_bilstm, InteractionNetwork_bilstm_full
from tqdm import tqdm
import argparse
def train(network, criterion, optimizer, dataloaders, stat_r, stat_v, n_epoch, augment_ratio,case_name, case_no):
os.system('mkdir -p ' + os.getcwd() +'/'+case_name+'/'+case_no+'/')
phases = ['train','valid']
writer = SummaryWriter(os.getcwd() +'/'+case_name+'/'+case_no+"/log") # use tensorboard
losses_epoch = {'train':[],'valid':[]}
for epoch in range(n_epoch):
start_epoch_time = time.time()
losses = {'train':[],'valid':[]}
for phase in phases:
network.train(phase == 'train')
loader = dataloaders[phase]
last_record_time = time.time()
last_record_time = time.time()
for i, data in enumerate(loader):
cnt_batch = 1+i
batch_data, label_data = data
batch_size = batch_data.size()[0]
batch_data = batch_data.cuda()
label_data = label_data.cuda()
if augment_ratio!=0:
# genearete noise to corrupt data
noise = torch.randn(batch_size,batch_data.size()[1]-4, 4).cuda() * stat_v[1,4:] * augment_ratio
batch_data[:,2:-2,:4] += noise
batch_data[:,2:-2,4:] += noise
object_for_relation_net, object_for_object_net, sender_relations, receiver_relations, target = prepare_data_pt(batch_size, batch_data, label_data, n_objects, n_relations)
predicted = network(object_for_relation_net,object_for_object_net, sender_relations, receiver_relations, stat_r,stat_v, 0)
target_n = (target-stat_v[0,4:])/stat_v[1,4:]
loss = criterion(predicted[2:-2], target_n[2:-2])
if phase == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses[phase].append(loss.data.item())
if cnt_batch % 100 == 0:
cost_time = time.time() - last_record_time
print('Epoch %s, Batch %s, RMS Error %s, time %s' % (epoch, cnt_batch, np.mean(losses[phase][-100:]), cost_time))
last_record_time = time.time()
losses_epoch[phase].append(np.mean(losses[phase]))
writer.add_scalar('Loss_mean/'+phase, losses_epoch[phase][-1], epoch+1)
parameter_file_path = os.getcwd() +'/'+case_name+'/'+case_no+"/epoch_"+str(epoch)+".pth"
torch.save(network.state_dict(), parameter_file_path)
print('Epoch costs %s(s)' % (time.time() - start_epoch_time))
# --- save state of both model and optimizer ---
parameter_file_path = os.getcwd() +'/'+case_name+'/'+case_no+"/t_epoch_"+str(epoch)+".pth"
state = {'model': network.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
torch.save(state, parameter_file_path)
if __name__ == '__main__':
# --- load data ---
parser = argparse.ArgumentParser()
parser.add_argument('--pair_wise_type', default='rel', help='relative state| full state concatenate, pair-wise relation type') # 'rel' or 'full'
parser.add_argument('--num_rnnlayer', type=int, default=2, help="the number of bilstm layers")
args = parser.parse_args()
timestep_per_scenario = 300
timestep_per_scenario_train = timestep_per_scenario-1
train_scenario = 2700
DATA_PATH = os.getcwd()+'/data/'
data_file = DATA_PATH + 'data_pt.npy'
data = np.load(data_file)
train_data = data[0:train_scenario]
valid_data = data[train_scenario:]
print('training_data', train_data.shape)
print('valid_data',valid_data.shape)
stat_r = np.load(DATA_PATH+"relation_vector_stat_pt.npy") # stat of relative position, relative twist, relative delta pos, relative delta twist
stat_v = np.load(DATA_PATH+"v_stat_pt.npy") # stat of position, twist, delta pos, delta twist
#################
CASE_NAME = 'test'
lr = 0.0001
batch_size = 512
epoch = 500
reps = 1
n_objects = 16 + 1
n_relations = (n_objects-1)*2
# test_training_data_normalized = np.array([training_data_normalized[:,:,:,:]])
import random
# random.seed(10)
augment_ratio = 0.05
num_layers = args.num_rnnlayer # 1,2,3
hidden_size = 150
if args.pair_wise_type == 'rel':
object_net_dim = 4+2 # [dx,dy,dz, dtwist, attri(2)]
relation_net_dim = 8+2*2 # [x,y,z,twist, dx,dy,dz,dtwist, attri(2)]
network_chosen = InteractionNetwork_bilstm
elif args.pair_wise_type == 'full':
object_net_dim = 4+2 # [dx,dy,dz, dtwist, attri(2)]
relation_net_dim = 8*2+2*2 # [x,y,z,twist, dx,dy,dz,dtwist, attri(2)]*2
network_chosen = InteractionNetwork_bilstm_full
stat_r = stat_v
else:
raise AssertionError("Unsupported pair-wise relation type, either relative ('rel') or full ('full')")
datasets=dict()
datasets['train'] = RopeDataset(train_data)
datasets['valid'] = RopeDataset(valid_data)
print("all data len, train:",datasets['train'].len, "valid:",datasets['valid'].len)
dataloaders = {x:DataLoader(
datasets[x], batch_size=batch_size,
shuffle=True if x == 'train' else False,
num_workers=8)
for x in ['train', 'valid']}
for k in range(reps):
network = network_chosen(relation_net_dim, object_net_dim, hidden_size, num_layers)
criterion = nn.MSELoss()
# criterion = nn.L1Loss()
network = network.cuda()
criterion = criterion.cuda()
optimizer = optim.Adam(network.parameters(), lr=lr, betas=(0.9, 0.999))
stat_r_ts = torch.Tensor(stat_r).cuda()
stat_v_ts = torch.Tensor(stat_v).cuda()
train(network, criterion, optimizer, dataloaders,stat_r_ts,stat_v_ts, epoch, augment_ratio, CASE_NAME, str(k))
import numpy as np
from torch.autograd import Variable
import torch
def prepare_data_pt(batch_size, batch_data, label_data, n_objects, n_relations):
#receiver_relations, sender_relations - onehot encoding matrices
#each column indicates the receiver and sender object’s index
receiver_relations = np.zeros((batch_size, n_objects, n_relations), dtype=float);
sender_relations = np.zeros((batch_size, n_objects, n_relations), dtype=float);
object_attr = torch.zeros((batch_size, n_objects, 2)).type('torch.FloatTensor').cuda()
object_attr[:,0:2,0] = torch.ones_like(object_attr[:,0:2,0]).type('torch.FloatTensor').cuda()
object_attr[:,2:-2,1] = torch.ones_like(object_attr[:,2:-2,1]).type('torch.FloatTensor').cuda()
object_attr[:,-2:,0] = torch.ones_like(object_attr[:,-2:,1]).type('torch.FloatTensor').cuda()
cnt = 0
for i in range(n_objects):
if i == n_objects-1:
receiver_relations[:, i, cnt] = 1.0
sender_relations[:, i - 1, cnt] = 1.0
cnt += 1
elif i == 0:
receiver_relations[:, i, cnt] = 1.0
sender_relations[:, i+1, cnt] = 1.0
cnt += 1
elif i != 0:
receiver_relations[:, i, cnt] = 1.0
sender_relations[:, i-1, cnt] = 1.0
cnt += 1
receiver_relations[:, i, cnt] = 1.0
sender_relations[:, i+1, cnt] = 1.0
cnt += 1
#There is no relation info in solar system task, just fill with zeros
target = label_data[:,:,4:]
object_for_relation_net = torch.cat([batch_data[:,:,:],object_attr],-1)
object_for_relation_net[:,-2:,:-2] = label_data[:,-2:,:]
object_for_object_net = object_for_relation_net
sender_relations = torch.FloatTensor(sender_relations)
receiver_relations = torch.FloatTensor(receiver_relations)
target = target.reshape(-1, 4)
sender_relations = sender_relations.cuda()
receiver_relations = receiver_relations.cuda()
return object_for_relation_net, object_for_object_net, sender_relations, receiver_relations, target
\ No newline at end of file