From c8737dc2667b4a5dbdbe5c514235d89ce9709d30 Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Mon, 28 Oct 2019 22:57:42 +0200 Subject: [PATCH] Fix Windows support in example code --- examples/pytorch_tensorboard.py | 113 +++++++++++++++++--------------- 1 file changed, 59 insertions(+), 54 deletions(-) diff --git a/examples/pytorch_tensorboard.py b/examples/pytorch_tensorboard.py index 9a85b168..5e7efcd4 100644 --- a/examples/pytorch_tensorboard.py +++ b/examples/pytorch_tensorboard.py @@ -15,50 +15,10 @@ from torch.autograd import Variable from torch.utils.tensorboard import SummaryWriter from trains import Task -task = Task.init(project_name='examples', task_name='pytorch with tensorboard') - - -writer = SummaryWriter('runs') -writer.add_text('lstm', 'This is an lstm', 0) -# Training settings -parser = argparse.ArgumentParser(description='PyTorch MNIST Example') -parser.add_argument('--batch-size', type=int, default=64, metavar='N', - help='input batch size for training (default: 64)') -parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', - help='input batch size for testing (default: 1000)') -parser.add_argument('--epochs', type=int, default=2, metavar='N', - help='number of epochs to train (default: 10)') -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('--no-cuda', action='store_true', default=False, - help='disables 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.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': 4, 'pin_memory': True} if args.cuda else {} -train_loader = torch.utils.data.DataLoader(datasets.MNIST('../data', 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', train=False, - transform=transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize((0.1307,), (0.3081,))])), - batch_size=args.batch_size, shuffle=True, **kwargs) class Net(nn.Module): + def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 10, kernel_size=5) @@ -77,13 +37,7 @@ class Net(nn.Module): return F.log_softmax(x) -model = Net() -if args.cuda: - model.cuda() -optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) - - -def train(epoch): +def internal_train(model, epoch, train_loader, args, optimizer, writer): model.train() for batch_idx, (data, target) in enumerate(train_loader): if args.cuda: @@ -102,7 +56,7 @@ def train(epoch): writer.add_scalar('Train/Loss', loss.data.item(), niter) -def test(): +def test(model, test_loader, args, optimizer, writer): model.eval() test_loss = 0 correct = 0 @@ -111,7 +65,7 @@ def test(): data, target = data.cuda(), target.cuda() data, target = Variable(data, volatile=True), Variable(target) output = model(data) - test_loss += F.nll_loss(output, target, size_average=False).data.item() # sum up batch loss + test_loss += F.nll_loss(output, target, size_average=False).data.item() # sum up batch loss pred = output.data.max(1)[1] # get the index of the max log-probability pred = pred.eq(target.data).cpu().sum() writer.add_scalar('Test/Loss', pred, niter) @@ -123,7 +77,58 @@ def test(): 100. * correct / len(test_loader.dataset))) -for epoch in range(1, args.epochs + 1): - train(epoch) - torch.save(model, os.path.join(gettempdir(), 'model{}'.format(epoch))) -test() +def main(): + # Training settings + parser = argparse.ArgumentParser(description='PyTorch MNIST Example') + parser.add_argument('--batch-size', type=int, default=64, metavar='N', + help='input batch size for training (default: 64)') + parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', + help='input batch size for testing (default: 1000)') + parser.add_argument('--epochs', type=int, default=2, metavar='N', + help='number of epochs to train (default: 10)') + 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('--no-cuda', action='store_true', default=False, + help='disables 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() + task = Task.init(project_name='examples', task_name='pytorch with tensorboard') + writer = SummaryWriter('runs') + writer.add_text('lstm', 'This is an lstm', 0) + 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': 4, 'pin_memory': True} if args.cuda else {} + train_loader = torch.utils.data.DataLoader(datasets.MNIST('../data', 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', train=False, + transform=transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.1307,), (0.3081,))])), + batch_size=args.batch_size, shuffle=True, **kwargs) + + model = Net() + if args.cuda: + model.cuda() + optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) + + for epoch in range(1, args.epochs + 1): + internal_train(model, epoch, train_loader, args, optimizer, writer) + torch.save(model, os.path.join(gettempdir(), 'model{}'.format(epoch))) + test(model, test_loader, args, optimizer, writer) + + +if __name__ == "__main__": + # Hack for supporting Windows OS - https://pytorch.org/docs/stable/notes/windows.html#usage-multiprocessing + main()