96 lines
2.7 KiB
Python
96 lines
2.7 KiB
Python
|
import os
|
||
|
|
import numpy as np
|
||
|
|
from sklearn.linear_model import LogisticRegression
|
||
|
|
import torch
|
||
|
|
import torchvision.datasets
|
||
|
|
import torchvision.models
|
||
|
|
import torchvision.transforms
|
||
|
|
|
||
|
|
from rbm import RBM
|
||
|
|
|
||
|
|
|
||
|
|
########## CONFIGURATION ##########
|
||
|
|
BATCH_SIZE = 64
|
||
|
|
VISIBLE_UNITS = 784 # 28 x 28 images
|
||
|
|
HIDDEN_UNITS = 128
|
||
|
|
CD_K = 2
|
||
|
|
EPOCHS = 100
|
||
|
|
|
||
|
|
DATA_FOLDER = os.path.split(os.path.realpath(__file__))[0]+'/../Dataset/'
|
||
|
|
print("Dataset Path :" + DATA_FOLDER)
|
||
|
|
|
||
|
|
CUDA = torch.cuda.is_available()
|
||
|
|
CUDA_DEVICE = 0
|
||
|
|
|
||
|
|
if CUDA:
|
||
|
|
torch.cuda.set_device(CUDA_DEVICE)
|
||
|
|
|
||
|
|
|
||
|
|
########## LOADING DATASET ##########
|
||
|
|
print('Loading dataset...')
|
||
|
|
|
||
|
|
train_dataset = torchvision.datasets.MNIST(root=DATA_FOLDER, train=True, transform=torchvision.transforms.ToTensor(), download=True)
|
||
|
|
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE)
|
||
|
|
|
||
|
|
test_dataset = torchvision.datasets.MNIST(root=DATA_FOLDER, train=False, transform=torchvision.transforms.ToTensor(), download=True)
|
||
|
|
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE)
|
||
|
|
|
||
|
|
|
||
|
|
########## TRAINING RBM ##########
|
||
|
|
print('Training RBM...')
|
||
|
|
|
||
|
|
rbm = RBM(VISIBLE_UNITS, HIDDEN_UNITS, CD_K, use_cuda=CUDA)
|
||
|
|
|
||
|
|
for epoch in range(EPOCHS):
|
||
|
|
epoch_error = 0.0
|
||
|
|
|
||
|
|
for batch, _ in train_loader:
|
||
|
|
batch = batch.view(len(batch), VISIBLE_UNITS) # flatten input data
|
||
|
|
|
||
|
|
if CUDA:
|
||
|
|
batch = batch.cuda()
|
||
|
|
|
||
|
|
batch_error = rbm.contrastive_divergence(batch)
|
||
|
|
|
||
|
|
epoch_error += batch_error
|
||
|
|
|
||
|
|
print('Epoch Error (epoch=%d): %.4f' % (epoch, epoch_error))
|
||
|
|
|
||
|
|
|
||
|
|
########## EXTRACT FEATURES ##########
|
||
|
|
print('Extracting features...')
|
||
|
|
|
||
|
|
train_features = np.zeros((len(train_dataset), HIDDEN_UNITS))
|
||
|
|
train_labels = np.zeros(len(train_dataset))
|
||
|
|
test_features = np.zeros((len(test_dataset), HIDDEN_UNITS))
|
||
|
|
test_labels = np.zeros(len(test_dataset))
|
||
|
|
|
||
|
|
for i, (batch, labels) in enumerate(train_loader):
|
||
|
|
batch = batch.view(len(batch), VISIBLE_UNITS) # flatten input data
|
||
|
|
|
||
|
|
if CUDA:
|
||
|
|
batch = batch.cuda()
|
||
|
|
|
||
|
|
train_features[i*BATCH_SIZE:i*BATCH_SIZE+len(batch)] = rbm.sample_hidden(batch).cpu().numpy()
|
||
|
|
train_labels[i*BATCH_SIZE:i*BATCH_SIZE+len(batch)] = labels.numpy()
|
||
|
|
|
||
|
|
for i, (batch, labels) in enumerate(test_loader):
|
||
|
|
batch = batch.view(len(batch), VISIBLE_UNITS) # flatten input data
|
||
|
|
|
||
|
|
if CUDA:
|
||
|
|
batch = batch.cuda()
|
||
|
|
|
||
|
|
test_features[i*BATCH_SIZE:i*BATCH_SIZE+len(batch)] = rbm.sample_hidden(batch).cpu().numpy()
|
||
|
|
test_labels[i*BATCH_SIZE:i*BATCH_SIZE+len(batch)] = labels.numpy()
|
||
|
|
|
||
|
|
|
||
|
|
########## CLASSIFICATION ##########
|
||
|
|
print('Classifying...')
|
||
|
|
|
||
|
|
clf = LogisticRegression()
|
||
|
|
clf.fit(train_features, train_labels)
|
||
|
|
predictions = clf.predict(test_features)
|
||
|
|
|
||
|
|
print('Result: %d/%d' % (sum(predictions == test_labels), test_labels.shape[0]))
|
||
|
|
|