Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable run main without argparse #19

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Enable run main without argparse #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4486691
correct error of not using cuda when CUDA is available.
Roller44 Jul 11, 2023
a39c245
correct error shape of transform for mnist dataset.
Roller44 Jul 11, 2023
9e27a97
Speed up distance calculation.
Roller44 Jul 12, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 2 additions & 6 deletions functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,9 @@ def forward(ctx, inputs, codebook):
embedding_size = codebook.size(1)
inputs_size = inputs.size()
inputs_flatten = inputs.view(-1, embedding_size)

codebook_sqr = torch.sum(codebook ** 2, dim=1)
inputs_sqr = torch.sum(inputs_flatten ** 2, dim=1, keepdim=True)


# Compute the distances to the codebook
distances = torch.addmm(codebook_sqr + inputs_sqr,
inputs_flatten, codebook.t(), alpha=-2.0, beta=1.0)
distances = torch.cdist(inputs_flatten, codebook, 2)

_, indices_flatten = torch.min(distances, dim=1)
indices = indices_flatten.view(*inputs_size[:-1])
Expand Down
11 changes: 5 additions & 6 deletions vqvae.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ def main(args):
if args.dataset in ['mnist', 'fashion-mnist', 'cifar10']:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
transforms.Normalize((0.5), (0.5))
])
if args.dataset == 'mnist':
# Define the train & test datasets
Expand Down Expand Up @@ -150,9 +150,9 @@ def main(args):
parser = argparse.ArgumentParser(description='VQ-VAE')

# General
parser.add_argument('--data-folder', type=str,
parser.add_argument('--data-folder', type=str, default='./data',
help='name of the data folder')
parser.add_argument('--dataset', type=str,
parser.add_argument('--dataset', type=str, default='mnist',
help='name of the dataset (mnist, fashion-mnist, cifar10, miniimagenet)')

# Latent space
Expand All @@ -176,7 +176,7 @@ def main(args):
help='name of the output folder (default: vqvae)')
parser.add_argument('--num-workers', type=int, default=mp.cpu_count() - 1,
help='number of workers for trajectories sampling (default: {0})'.format(mp.cpu_count() - 1))
parser.add_argument('--device', type=str, default='cpu',
parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
help='set the device (cpu or cuda, default: cpu)')

args = parser.parse_args()
Expand All @@ -187,8 +187,7 @@ def main(args):
if not os.path.exists('./models'):
os.makedirs('./models')
# Device
args.device = torch.device(args.device
if torch.cuda.is_available() else 'cpu')
args.device = torch.device(args.device)
# Slurm
if 'SLURM_JOB_ID' in os.environ:
args.output_folder += '-{0}'.format(os.environ['SLURM_JOB_ID'])
Expand Down