Fix tfrecord-converter

Author: wohlert

run-gqn.py

@@ -67,6 +67,8 @@
     valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)
 
     def step(engine, batch):
+        model.train()
+
         x, v = batch
         x, v = x.to(device), v.to(device)
         x, v, x_q, v_q = partition(x, v)
@@ -101,7 +103,10 @@ def step(engine, batch):
     # Trainer and metrics
     trainer = Engine(step)
     metric_names = ["elbo", "kl", "sigma", "mu"]
-    metrics = [RunningAverage(output_transform=lambda x: x[m]).attach(trainer, m) for m in metric_names]
+    RunningAverage(output_transform=lambda x: x["elbo"]).attach(trainer, "elbo")
+    RunningAverage(output_transform=lambda x: x["kl"]).attach(trainer, "kl")
+    RunningAverage(output_transform=lambda x: x["sigma"]).attach(trainer, "sigma")
+    RunningAverage(output_transform=lambda x: x["mu"]).attach(trainer, "mu")
     ProgressBar().attach(trainer, metric_names=metric_names)
 
     # Model checkpointing
@@ -142,6 +147,7 @@ def save_images(engine):
 
     @trainer.on(Events.EPOCH_COMPLETED)
     def validate(engine):
+        model.eval()
         with torch.no_grad():
             x, v = next(iter(valid_loader))
             x, v = x.to(device), v.to(device)

scripts/data.sh

@@ -1,17 +1,20 @@
 #!/usr/bin/env bash
 
-LOCATION=$1
-BATCH_SIZE=$2
+LOCATION=$1   # example: /tmp/data
+BATCH_SIZE=$2 # example: 64
 
 echo "Downloading data"
 gsutil -m cp -R gs://gqn-dataset/shepard_metzler_5_parts $LOCATION
 
-echo "Deleting small records"
-TRAIN_PATH="$LOCATION/shepard_metzler_5_parts/train"
-find "$TRAIN_PATH/*.tfrecord" -type f -size -10M | xargs rm # remove smaller than 10mb
+echo "Deleting small records" # less than 10MB
+DATA_PATH="$LOCATION/shepard_metzler_5_parts/**/*.tfrecord"
+find $DATA_PATH -type f -size -10M | xargs rm
 
 echo "Converting data"
 python tfrecord-converter.py $LOCATION shepard_metzler_5_parts -b $BATCH_SIZE -m "train"
 echo "Training data: done"
 python tfrecord-converter.py $LOCATION shepard_metzler_5_parts -b $BATCH_SIZE -m "test"
-echo "Testing data: done"
+echo "Testing data: done"
+
+echo "Removing original records"
+rm -rf "$LOCATION/shepard_metzler_5_parts/**/*.tfrecord"

scripts/gpu.sh

@@ -14,4 +14,4 @@ python ../run-gqn.py \
     --log_dir "../logs" \
     --data_parallel "True" \
     --batch_size 1 \
-    --n_workers 6
+    --workers 6

scripts/tfrecord-converter.py

@@ -46,7 +46,7 @@ def process(record):
 
         # Convert
         images = tf.map_fn(tf.image.decode_jpeg, tf.reshape(images, [-1]), **kwargs)
-        images = tf.reshape(images, (-1, SEQ_DIM, 3, IMG_DIM, IMG_DIM))
+        images = tf.reshape(images, (-1, SEQ_DIM, IMG_DIM, IMG_DIM, 3))
         poses  = tf.reshape(poses,  (-1, SEQ_DIM, POSE_DIM))
 
         # Numpy conversion
@@ -64,8 +64,8 @@ def convert(record, batch_size):
     batch_process = lambda r: chunk(process(r), batch_size)
 
     for i, batch in enumerate(batch_process(record)):
-        path = os.path.join(path, "{0:}-{1:02}.pt.gz".format(basename, i))
-        with gzip.open(path, 'wb') as f:
+        p = os.path.join(path, "{0:}-{1:02}.pt.gz".format(basename, i))
+        with gzip.open(p, 'wb') as f:
             torch.save(list(batch), f)
 
 if __name__ == '__main__':
@@ -91,4 +91,4 @@ def convert(record, batch_size):
 
     with mp.Pool(processes=mp.cpu_count()) as pool:
         f = partial(convert, batch_size=args.batch_size)
-        pool.map(f, records)
+        pool.map(f, records)

shepardmetzler.py

@@ -49,7 +49,8 @@ def __getitem__(self, idx):
         images, viewpoints = list(zip(*data))
 
         # (b, m, c, h, w)
-        images = torch.FloatTensor(images)
+        images = torch.FloatTensor(images)/255
+        images = images.permute(0, 1, 4, 2, 3)
         if self.transform:
             images = self.transform(images)