246 lines
8.9 KiB
Python
246 lines
8.9 KiB
Python
|
"""
|
||
|
|
A deep MNIST classifier using convolutional layers.
|
||
|
|
|
||
|
|
This file is a modification of the official pytorch mnist example:
|
||
|
|
https://github.com/pytorch/examples/blob/master/mnist/main.py
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
import argparse
|
||
|
|
import logging
|
||
|
|
import nni
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
import torch.optim as optim
|
||
|
|
from torchvision import datasets, transforms
|
||
|
|
|
||
|
|
from dataset import *
|
||
|
|
from models import *
|
||
|
|
|
||
|
|
logger = logging.getLogger('IMPAX_AutoML')
|
||
|
|
|
||
|
|
criterion = nn.MSELoss()
|
||
|
|
# criterion = nn.MSELoss(reduction='sum')
|
||
|
|
# criterion = nn.MSELoss
|
||
|
|
|
||
|
|
# MODEL_DIR = os.path.dirname(os.path.realpath(__file__))
|
||
|
|
MODEL_DIR = "/home/xfr/nni/"
|
||
|
|
|
||
|
|
def train(args, model, device, train_loader, optimizer, epoch):
|
||
|
|
|
||
|
|
model.train()
|
||
|
|
train_loss = 0
|
||
|
|
for batch_idx, (data, target) in enumerate(train_loader):
|
||
|
|
data, target = data.to(device), target.to(device)
|
||
|
|
optimizer.zero_grad()
|
||
|
|
output = model(data)
|
||
|
|
loss = criterion(output, target)
|
||
|
|
train_loss += loss * len(data)
|
||
|
|
loss.backward()
|
||
|
|
optimizer.step()
|
||
|
|
if batch_idx % args['log_interval'] == 0:
|
||
|
|
logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
|
||
|
|
epoch, batch_idx * len(data), len(train_loader.dataset),
|
||
|
|
100. * batch_idx / len(train_loader), loss.item()))
|
||
|
|
logger.info('Train Epoch {}:\tAverage Loss: {:.6f}'.format(
|
||
|
|
epoch,
|
||
|
|
train_loss / len(train_loader.dataset),
|
||
|
|
))
|
||
|
|
|
||
|
|
# nni.get_experiment_id()
|
||
|
|
# nni.get_trial_id()
|
||
|
|
# nni.get_sequence_id()
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def test(args, model, device, test_loader):
|
||
|
|
model.eval()
|
||
|
|
test_loss = 0
|
||
|
|
correct = 0
|
||
|
|
with torch.no_grad():
|
||
|
|
for data, target in test_loader:
|
||
|
|
data, target = data.to(device), target.to(device)
|
||
|
|
output = model(data)
|
||
|
|
# sum up batch loss
|
||
|
|
# logger.info('criterion(output, target).item() %s' % criterion(output, target).item())
|
||
|
|
# logger.info('len(test_loader) %s' % len(test_loader))
|
||
|
|
# logger.info('len(data) %s' % len(data))
|
||
|
|
test_loss += criterion(output, target).item() * len(data)
|
||
|
|
# get the index of the max log-probability
|
||
|
|
pred = output.argmax(dim=1, keepdim=True)
|
||
|
|
# correct += pred.eq(target.view_as(pred)).sum().item()
|
||
|
|
|
||
|
|
test_loss /= len(test_loader.dataset)
|
||
|
|
|
||
|
|
# accuracy = 100. * correct / len(test_loader.dataset)
|
||
|
|
|
||
|
|
# logger.info('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
|
||
|
|
# test_loss, correct, len(test_loader.dataset), accuracy))
|
||
|
|
|
||
|
|
logger.info('Test set: Average loss: {:.4f}, {}'.format(
|
||
|
|
test_loss,
|
||
|
|
len(test_loader.dataset),
|
||
|
|
))
|
||
|
|
|
||
|
|
|
||
|
|
return test_loss
|
||
|
|
|
||
|
|
|
||
|
|
# def main(args):
|
||
|
|
# use_cuda = not args['no_cuda'] and torch.cuda.is_available()
|
||
|
|
|
||
|
|
# torch.manual_seed(args['seed'])
|
||
|
|
|
||
|
|
# device = torch.device("cuda" if use_cuda else "cpu")
|
||
|
|
|
||
|
|
# kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
|
||
|
|
|
||
|
|
# data_dir = os.path.join(args['data_dir'], nni.get_trial_id())
|
||
|
|
|
||
|
|
# train_loader = torch.utils.data.DataLoader(
|
||
|
|
# datasets.MNIST(data_dir, train=True, download=True,
|
||
|
|
# transform=transforms.Compose([
|
||
|
|
# transforms.ToTensor(),
|
||
|
|
# transforms.Normalize((0.1307,), (0.3081,))
|
||
|
|
# ])),
|
||
|
|
# batch_size=args['batch_size'], shuffle=True, **kwargs)
|
||
|
|
# test_loader = torch.utils.data.DataLoader(
|
||
|
|
# datasets.MNIST(data_dir, train=False, transform=transforms.Compose([
|
||
|
|
# transforms.ToTensor(),
|
||
|
|
# transforms.Normalize((0.1307,), (0.3081,))
|
||
|
|
# ])),
|
||
|
|
# batch_size=1000, shuffle=True, **kwargs)
|
||
|
|
|
||
|
|
# hidden_size = args['hidden_size']
|
||
|
|
|
||
|
|
# model = Net(hidden_size=hidden_size).to(device)
|
||
|
|
# optimizer = optim.SGD(model.parameters(), lr=args['lr'],
|
||
|
|
# momentum=args['momentum'])
|
||
|
|
|
||
|
|
# for epoch in range(1, args['epochs'] + 1):
|
||
|
|
# train(args, model, device, train_loader, optimizer, epoch)
|
||
|
|
# test_acc = test(args, model, device, test_loader)
|
||
|
|
|
||
|
|
# if epoch < args['epochs']:
|
||
|
|
# # report intermediate result
|
||
|
|
# nni.report_intermediate_result(test_acc)
|
||
|
|
# logger.debug('test accuracy %g', test_acc)
|
||
|
|
# logger.debug('Pipe send intermediate result done.')
|
||
|
|
# else:
|
||
|
|
# # report final result
|
||
|
|
# nni.report_final_result(test_acc)
|
||
|
|
# logger.debug('Final result is %g', test_acc)
|
||
|
|
# logger.debug('Send final result done.')
|
||
|
|
|
||
|
|
|
||
|
|
def main(args):
|
||
|
|
use_cuda = not args['no_cuda'] and torch.cuda.is_available()
|
||
|
|
|
||
|
|
torch.manual_seed(args['seed'])
|
||
|
|
|
||
|
|
device = torch.device("cuda" if use_cuda else "cpu")
|
||
|
|
|
||
|
|
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
|
||
|
|
|
||
|
|
# data_dir = os.path.join(args['data_dir'], nni.get_trial_id())
|
||
|
|
trainset = IMPAXDataset(os.path.join(args['data_dir'], 'train'))
|
||
|
|
testset = IMPAXDataset(os.path.join(args['data_dir'], 'test'))
|
||
|
|
|
||
|
|
train_loader = torch.utils.data.DataLoader(trainset, batch_size=args['batch_size'],
|
||
|
|
shuffle=True, **kwargs)
|
||
|
|
|
||
|
|
test_loader = torch.utils.data.DataLoader(testset,
|
||
|
|
batch_size=100,
|
||
|
|
shuffle=True, **kwargs)
|
||
|
|
|
||
|
|
hidden_layer = args['hidden_layer']
|
||
|
|
hidden_size = args['hidden_size']
|
||
|
|
|
||
|
|
model = Net(hidden_layer=hidden_layer, hidden_size=hidden_size).to(device)
|
||
|
|
optimizer = optim.Adam(model.parameters(),
|
||
|
|
lr=args['lr'],
|
||
|
|
# momentum=args['momentum'],
|
||
|
|
)
|
||
|
|
|
||
|
|
best_loss = None
|
||
|
|
for epoch in range(1, args['epochs'] + 1):
|
||
|
|
train(args, model, device, train_loader, optimizer, epoch)
|
||
|
|
test_loss = test(args, model, device, test_loader)
|
||
|
|
|
||
|
|
if best_loss is None or best_loss > test_loss :
|
||
|
|
best_loss = test_loss
|
||
|
|
|
||
|
|
model_subdir = nni.get_experiment_id()
|
||
|
|
|
||
|
|
if args['exp_name'] is None:
|
||
|
|
model_file = os.path.join(MODEL_DIR, model.name, model_subdir, 'best_{}.pth'.format(nni.get_trial_id()))
|
||
|
|
else:
|
||
|
|
model_file = os.path.join(MODEL_DIR, args['exp_name'], model.name, model_subdir, 'best_{}.pth'.format(nni.get_trial_id()))
|
||
|
|
|
||
|
|
parent_dir = os.path.dirname(model_file)
|
||
|
|
if not os.path.exists(parent_dir):
|
||
|
|
os.makedirs(parent_dir)
|
||
|
|
torch.save(model.state_dict(), model_file)
|
||
|
|
|
||
|
|
logger.info('model saved: %s' % model_file)
|
||
|
|
|
||
|
|
|
||
|
|
if epoch < args['epochs']:
|
||
|
|
# report intermediate result
|
||
|
|
nni.report_intermediate_result(test_loss)
|
||
|
|
logger.debug('test loss %g', test_loss)
|
||
|
|
logger.debug('Pipe send intermediate result done.')
|
||
|
|
else:
|
||
|
|
# report final result
|
||
|
|
nni.report_final_result(test_loss)
|
||
|
|
logger.debug('Final result is %g', test_loss)
|
||
|
|
logger.debug('Send final result done.')
|
||
|
|
|
||
|
|
logger.info(' ')
|
||
|
|
|
||
|
|
|
||
|
|
def get_params():
|
||
|
|
# Training settings
|
||
|
|
parser = argparse.ArgumentParser(description='PyTorch IMPAX Example')
|
||
|
|
|
||
|
|
# parser.add_argument("--data_dir", type=str,
|
||
|
|
# default='/tmp/tensorflow/mnist/input_data', help="data directory")
|
||
|
|
|
||
|
|
parser.add_argument("--data_dir", type=str,
|
||
|
|
default='/shares/Public/IMPAX/', help="data directory")
|
||
|
|
|
||
|
|
parser.add_argument('--batch_size', type=int, default=64, metavar='N',
|
||
|
|
help='input batch size for training (default: 64)')
|
||
|
|
parser.add_argument("--hidden_size", type=int, default=512, metavar='N',
|
||
|
|
help='hidden layer size (default: 512)')
|
||
|
|
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
|
||
|
|
help='learning rate (default: 0.01)')
|
||
|
|
parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
|
||
|
|
help='SGD momentum (default: 0.5)')
|
||
|
|
parser.add_argument('--epochs', type=int, default=10, metavar='N',
|
||
|
|
help='number of epochs to train (default: 10)')
|
||
|
|
parser.add_argument('--seed', type=int, default=1, metavar='S',
|
||
|
|
help='random seed (default: 1)')
|
||
|
|
parser.add_argument('--no_cuda', action='store_true', default=False,
|
||
|
|
help='disables CUDA training')
|
||
|
|
parser.add_argument('--log_interval', type=int, default=1000, metavar='N',
|
||
|
|
help='how many batches to wait before logging training status')
|
||
|
|
|
||
|
|
parser.add_argument('--exp_name', default=None, type=str, help='exp name')
|
||
|
|
|
||
|
|
args, _ = parser.parse_known_args()
|
||
|
|
return args
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
try:
|
||
|
|
# get parameters form tuner
|
||
|
|
tuner_params = nni.get_next_parameter()
|
||
|
|
logger.debug(tuner_params)
|
||
|
|
params = vars(get_params())
|
||
|
|
params.update(tuner_params)
|
||
|
|
main(params)
|
||
|
|
except Exception as exception:
|
||
|
|
logger.exception(exception)
|
||
|
|
raise
|