nt = 10 # num of time steps
layer_loss_weights = Variable(
torch.FloatTensor([[1.], [0.], [0.], [0.]]).cuda())
time_loss_weights = 1. / (nt - 1) * torch.ones(nt, 1)
time_loss_weights[0] = 0
time_loss_weights = Variable(time_loss_weights.cuda())
DATA_DIR = 'kitti_data_raw'
train_file = os.path.join(DATA_DIR, 'X_train.hkl')
train_sources = os.path.join(DATA_DIR, 'sources_train.hkl')
val_file = os.path.join(DATA_DIR, 'X_val.hkl')
val_sources = os.path.join(DATA_DIR, 'sources_val.hkl')
kitti_train = KITTI(train_file, train_sources, nt)
kitti_val = KITTI(val_file, val_sources, nt)
train_loader = DataLoader(kitti_train, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(kitti_val, batch_size=batch_size, shuffle=True)
model = PredNet(R_channels, A_channels, output_mode='error')
if torch.cuda.is_available():
print('Using GPU.')
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
def lr_scheduler(optimizer, epoch):
if epoch < num_epochs // 2:
from PIL import Image
im = Image.fromarray(np.rollaxis(tensor.numpy(), 0, 3))
im.save(filename)
from scipy.misc import imshow, imsave
batch_size = 16
A_channels = (3, 48, 96, 192)
R_channels = (3, 48, 96, 192)
DATA_DIR = '/media/lei/000F426D0004CCF4/datasets/kitti_data'
test_file = os.path.join(DATA_DIR, 'X_test.hkl')
test_sources = os.path.join(DATA_DIR, 'sources_test.hkl')
nt = 10
kitti_test = KITTI(test_file, test_sources, nt)
test_loader = DataLoader(kitti_test, batch_size=batch_size, shuffle=False)
model = PredNet(R_channels, A_channels, output_mode='prediction')
model.load_state_dict(torch.load('training.pt'))
if torch.cuda.is_available():
print('Using GPU.')
model.cuda()
for i, inputs in enumerate(test_loader):
inputs = inputs.permute(0, 1, 4, 2, 3) # batch x time_steps x channel x width x height
inputs = Variable(inputs.cuda())
origin = inputs.data.cpu().byte()[:, nt-1]
print('origin:')
nt = args.nt # num of time steps
if args.grayscale: # grayscale 1 channel
A_channels = (1, 24, 48, 96)
R_channels = (1, 24, 48, 96)
else: # RGB 3 channels
A_channels = (3, 48, 96, 192)
R_channels = (3, 48, 96, 192)
DATA_DIR = args.data_dir #'kitti_data'
TRAIN_DIR = 'trained'
RESULTS_DIR = 'results'
test_file = os.path.join(DATA_DIR, 'X_test.hkl')
test_sources = os.path.join(DATA_DIR, 'sources_test.hkl')
kitti_test = KITTI(test_file, test_sources, nt, sequence_start_mode='unique')
test_loader = DataLoader(kitti_test, batch_size=batch_size, shuffle=False)
model = PredNet(R_channels,
A_channels,
output_mode='prediction',
gpu_id=args.gpu_id)
model.load_state_dict(
torch.load(os.path.join(TRAIN_DIR, 'training.pt'),
map_location=lambda storage, loc: storage))
if args.et > 0 and args.et < args.nt:
model.set_extrap_start_time(extrap_start_time=args.et)
if args.gpu_id >= 0 and torch.cuda.is_available():
print(' Using GPU.')
[0.]])) #.cuda()) # L_0
#layer_loss_weights = Variable(torch.FloatTensor([[1.], [1.], [1.], [1.]])) #.cuda()) # L_all
time_loss_weights = 1. / (nt - 1) * torch.ones(nt, 1)
time_loss_weights[0] = 0
time_loss_weights = Variable(time_loss_weights) #.cuda())
DATA_DIR = args.data_dir #'kitti_data'
TRAIN_DIR = 'trained'
train_file = os.path.join(DATA_DIR, 'X_train.hkl')
train_sources = os.path.join(DATA_DIR, 'sources_train.hkl')
val_file = os.path.join(DATA_DIR, 'X_val.hkl')
val_sources = os.path.join(DATA_DIR, 'sources_val.hkl')
if args.gpu_id >= 0:
kitti_train = KITTI(train_file, train_sources, nt)
kitti_val = KITTI(val_file, val_sources, nt, N_seq=N_seq_val)
else:
kitti_train = KITTI(train_file, train_sources, nt, N_seq=500)
print(
'WARNNING: possible_starts is set to 500. Data volume is short for train.'
) # if not limit 500, memory will be over > 8G...
kitti_val = KITTI(val_file, val_sources, nt, N_seq=N_seq_val)
train_loader = DataLoader(kitti_train, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(kitti_val, batch_size=batch_size, shuffle=True)
model = PredNet(R_channels,
A_channels,
output_mode='error',
gpu_id=args.gpu_id)