-
Notifications
You must be signed in to change notification settings - Fork 3
/
main.py
90 lines (63 loc) · 2.73 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
from InfoCatVAE import InfoCatVAE
from utils import *
from train_test_functions import *
# ----- Parameters ----- #
parser = argparse.ArgumentParser(description='GMSVAE MNIST Example')
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=100, metavar='N',
help='number of epochs to train (default: 10)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
help='how many batches to wait before logging training status')
# args = parser.parse_args()
args, unknown = parser.parse_known_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
# ----- Data ----- #
train_data = datasets.MNIST('./data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
]))
test_data = datasets.MNIST('./data', train=False, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
]))
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size, shuffle=True, **kwargs)
# ----- Model parameters ----- #
loss_classifier = nn.NLLLoss()
n_epochs = 5000
batch_size = 64
in_dim = 784
num_class = 10
sub_dim = 2
z_dim = num_class*sub_dim
h_dim = 400
lmbda = 2
test_class_perf = []
test_lost_list = []
model = InfoCatVAE(in_dim, num_class, sub_dim, z_dim, h_dim, lmbda, args.cuda)
args.is_convolutional = False
# ----- Prior ----- #
args.mupriorT = prior_construction(model)
# ----- Learning ----- #
if args.cuda:
model.cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(1, n_epochs + 1):
train(epoch, model, train_loader, optimizer, args)
test(epoch, model, test_loader, args)
if epoch % 10 == 0:
torch.save(model.state_dict(), '../InfoCatVAE' + str(epoch) + '.pt')
sample, labels = sampling(model, 10, mupriorT)
save_image(sample.view(10*num_class, 1, 28, 28).data,
'../InfoCatVAE' + str(epoch) + '.png', nrow=num_class)