Commit 7aab770

committed

Fix tfrecord-converter

1 parent 08bf31e commit 7aab770Copy full SHA for 7aab770

File tree

+23

-13

lines changed

+23

-13

lines changed

Lines changed: 7 additions & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -67,6 +67,8 @@`
`67`	`67`	`valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)`
`68`	`68`
`69`	`69`	`def step(engine, batch):`
	`70`	`+ model.train()`
	`71`	`+`
`70`	`72`	`x, v = batch`
`71`	`73`	`x, v = x.to(device), v.to(device)`
`72`	`74`	`x, v, x_q, v_q = partition(x, v)`
`@@ -101,7 +103,10 @@ def step(engine, batch):`
`101`	`103`	`# Trainer and metrics`
`102`	`104`	`trainer = Engine(step)`
`103`	`105`	`metric_names = ["elbo", "kl", "sigma", "mu"]`
`104`		`- metrics = [RunningAverage(output_transform=lambda x: x[m]).attach(trainer, m) for m in metric_names]`
	`106`	`+ RunningAverage(output_transform=lambda x: x["elbo"]).attach(trainer, "elbo")`
	`107`	`+ RunningAverage(output_transform=lambda x: x["kl"]).attach(trainer, "kl")`
	`108`	`+ RunningAverage(output_transform=lambda x: x["sigma"]).attach(trainer, "sigma")`
	`109`	`+ RunningAverage(output_transform=lambda x: x["mu"]).attach(trainer, "mu")`
`105`	`110`	`ProgressBar().attach(trainer, metric_names=metric_names)`
`106`	`111`
`107`	`112`	`# Model checkpointing`
`@@ -142,6 +147,7 @@ def save_images(engine):`
`142`	`147`
`143`	`148`	`@trainer.on(Events.EPOCH_COMPLETED)`
`144`	`149`	`def validate(engine):`
	`150`	`+ model.eval()`
`145`	`151`	`with torch.no_grad():`
`146`	`152`	`x, v = next(iter(valid_loader))`
`147`	`153`	`x, v = x.to(device), v.to(device)`

Lines changed: 9 additions & 6 deletions

Original file line number	Diff line number	Diff line change
`@@ -1,17 +1,20 @@`
`1`	`1`	`#!/usr/bin/env bash`
`2`	`2`
`3`		`-LOCATION=1ドル`
`4`		`-BATCH_SIZE=2ドル`
	`3`	`+LOCATION=1ドル# example: /tmp/data`
	`4`	`+BATCH_SIZE=2ドル# example: 64`
`5`	`5`
`6`	`6`	`echo "Downloading data"`
`7`	`7`	`gsutil -m cp -R gs://gqn-dataset/shepard_metzler_5_parts $LOCATION`
`8`	`8`
`9`		`-echo "Deleting small records"`
`10`		`-TRAIN_PATH="$LOCATION/shepard_metzler_5_parts/train"`
`11`		`-find "$TRAIN_PATH/*.tfrecord" -type f -size -10M \| xargs rm# remove smaller than 10mb`
	`9`	`+echo "Deleting small records"# less than 10MB`
	`10`	`+DATA_PATH="$LOCATION/shepard_metzler_5_parts/*/.tfrecord"`
	`11`	`+find $DATA_PATH -type f -size -10M \| xargs rm`
`12`	`12`
`13`	`13`	`echo "Converting data"`
`14`	`14`	`python tfrecord-converter.py $LOCATION shepard_metzler_5_parts -b $BATCH_SIZE -m "train"`
`15`	`15`	`echo "Training data: done"`
`16`	`16`	`python tfrecord-converter.py $LOCATION shepard_metzler_5_parts -b $BATCH_SIZE -m "test"`
`17`		`-echo "Testing data: done"`
	`17`	`+echo "Testing data: done"`
	`18`	`+`
	`19`	`+echo "Removing original records"`
	`20`	`+rm -rf "$LOCATION/shepard_metzler_5_parts/*/.tfrecord"`

Lines changed: 1 addition & 1 deletion

Lines changed: 4 additions & 4 deletions

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ def process(record):`
`46`	`46`
`47`	`47`	`# Convert`
`48`	`48`	`images = tf.map_fn(tf.image.decode_jpeg, tf.reshape(images, [-1]), **kwargs)`
`49`		`- images = tf.reshape(images, (-1, SEQ_DIM, 3, IMG_DIM, IMG_DIM))`
	`49`	`+ images = tf.reshape(images, (-1, SEQ_DIM, IMG_DIM, IMG_DIM, 3))`
`50`	`50`	`poses = tf.reshape(poses, (-1, SEQ_DIM, POSE_DIM))`
`51`	`51`
`52`	`52`	`# Numpy conversion`
`@@ -64,8 +64,8 @@ def convert(record, batch_size):`
`64`	`64`	`batch_process = lambda r: chunk(process(r), batch_size)`
`65`	`65`
`66`	`66`	`for i, batch in enumerate(batch_process(record)):`
`67`		`- path = os.path.join(path, "{0:}-{1:02}.pt.gz".format(basename, i))`
`68`		`- with gzip.open(path, 'wb') as f:`
	`67`	`+ p = os.path.join(path, "{0:}-{1:02}.pt.gz".format(basename, i))`
	`68`	`+ with gzip.open(p, 'wb') as f:`
`69`	`69`	`torch.save(list(batch), f)`
`70`	`70`
`71`	`71`	`if __name__ == '__main__':`
`@@ -91,4 +91,4 @@ def convert(record, batch_size):`
`91`	`91`
`92`	`92`	`with mp.Pool(processes=mp.cpu_count()) as pool:`
`93`	`93`	`f = partial(convert, batch_size=args.batch_size)`
`94`		`- pool.map(f, records)`
	`94`	`+ pool.map(f, records)`

Lines changed: 2 additions & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,8 @@ def __getitem__(self, idx):`
`49`	`49`	`images, viewpoints = list(zip(*data))`
`50`	`50`
`51`	`51`	`# (b, m, c, h, w)`
`52`		`- images = torch.FloatTensor(images)`
	`52`	`+ images = torch.FloatTensor(images)/255`
	`53`	`+ images = images.permute(0, 1, 4, 2, 3)`
`53`	`54`	`if self.transform:`
`54`	`55`	`images = self.transform(images)`
`55`	`56`

Comments

(0)