""" 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