diff --git a/examples/pytorch_tensorboard.py b/examples/pytorch_tensorboard.py index 5e7efcd4..0e56f01a 100644 --- a/examples/pytorch_tensorboard.py +++ b/examples/pytorch_tensorboard.py @@ -34,10 +34,10 @@ class Net(nn.Module): x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) - return F.log_softmax(x) + return F.log_softmax(x, dim=1) -def internal_train(model, epoch, train_loader, args, optimizer, writer): +def train(model, epoch, train_loader, args, optimizer, writer): model.train() for batch_idx, (data, target) in enumerate(train_loader): if args.cuda: @@ -63,13 +63,15 @@ def test(model, test_loader, args, optimizer, writer): for niter, (data, target) in enumerate(test_loader): if args.cuda: data, target = data.cuda(), target.cuda() - data, target = Variable(data, volatile=True), Variable(target) + data, target = Variable(data), 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, reduction='sum').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) correct += pred + if niter % 100 == 0: + writer.add_image('test', data[0, :, :, :], niter) test_loss /= len(test_loader.dataset) print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( @@ -97,9 +99,9 @@ def main(): 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') + task = Task.init(project_name='examples', task_name='pytorch with tensorboard', output_uri='/tmp/blah') writer = SummaryWriter('runs') - writer.add_text('lstm', 'This is an lstm', 0) + writer.add_text('TEXT', 'This is some text', 0) args.cuda = not args.no_cuda and torch.cuda.is_available() torch.manual_seed(args.seed) @@ -124,7 +126,7 @@ def main(): 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) + 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) diff --git a/examples/pytorch_tensorboardX.py b/examples/pytorch_tensorboardX.py index 13133b53..51bfa77b 100644 --- a/examples/pytorch_tensorboardX.py +++ b/examples/pytorch_tensorboardX.py @@ -33,7 +33,49 @@ class Net(nn.Module): x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) - return F.log_softmax(x) + return F.log_softmax(x, dim=1) + + +def train(model, epoch, train_loader, args, optimizer, writer): + model.train() + for batch_idx, (data, target) in enumerate(train_loader): + if args.cuda: + data, target = data.cuda(), target.cuda() + data, target = Variable(data), Variable(target) + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % args.log_interval == 0: + print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + epoch, batch_idx * len(data), len(train_loader.dataset), + 100. * batch_idx / len(train_loader), loss.data.item())) + niter = epoch*len(train_loader)+batch_idx + writer.add_scalar('Train/Loss', loss.data.item(), niter) + + +def test(model, test_loader, args, optimizer, writer): + model.eval() + test_loss = 0 + correct = 0 + for niter, (data, target) in enumerate(test_loader): + if args.cuda: + data, target = data.cuda(), target.cuda() + data, target = Variable(data), Variable(target) + output = model(data) + test_loss += F.nll_loss(output, target, reduction='sum').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) + correct += pred + if niter % 100 == 0: + writer.add_image('test', data[0, :, :, :], niter) + + test_loss /= len(test_loader.dataset) + print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( + test_loss, correct, len(test_loader.dataset), + 100. * correct / len(test_loader.dataset))) def main(): @@ -60,7 +102,7 @@ def main(): task = Task.init(project_name='examples', task_name='pytorch with tensorboardX') writer = SummaryWriter('runs') - writer.add_text('lstm', 'This is an lstm', 0) + writer.add_text('TEXT', 'This is some text', 0) torch.manual_seed(args.seed) if args.cuda: @@ -83,48 +125,10 @@ def main(): model.cuda() optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) - def train(epoch): - model.train() - for batch_idx, (data, target) in enumerate(train_loader): - if args.cuda: - data, target = data.cuda(), target.cuda() - data, target = Variable(data), Variable(target) - optimizer.zero_grad() - output = model(data) - loss = F.nll_loss(output, target) - loss.backward() - optimizer.step() - if batch_idx % args.log_interval == 0: - print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( - epoch, batch_idx * len(data), len(train_loader.dataset), - 100. * batch_idx / len(train_loader), loss.data.item())) - niter = epoch*len(train_loader)+batch_idx - writer.add_scalar('Train/Loss', loss.data.item(), niter) - - def test(): - model.eval() - test_loss = 0 - correct = 0 - for niter, (data, target) in enumerate(test_loader): - if args.cuda: - 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 - 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) - correct += pred - - test_loss /= len(test_loader.dataset) - print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( - test_loss, correct, len(test_loader.dataset), - 100. * correct / len(test_loader.dataset))) - for epoch in range(1, args.epochs + 1): - train(epoch) + train(model, epoch, train_loader, args, optimizer, writer) torch.save(model, os.path.join(gettempdir(), 'model{}'.format(epoch))) - test() + test(model, test_loader, args, optimizer, writer) if __name__ == "__main__":