add flush to print (pytorch#81)

With flush, log info can appear immediately when it is directed to a pipe or file.

Author: fyu

imagenet/main.py

@@ -1,3 +1,4 @@
+from __future__ import print_function
 import argparse
 import os
 import shutil
@@ -59,10 +60,10 @@ def main():
 
  # create model
  if args.pretrained:
- print("=> using pre-trained model '{}'".format(args.arch))
+ print("=> using pre-trained model '{}'".format(args.arch), flush=True)
  model = models.__dict__[args.arch](pretrained=True)
  else:
- print("=> creating model '{}'".format(args.arch))
+ print("=> creating model '{}'".format(args.arch), flush=True)
  model = models.__dict__[args.arch]()
 
  if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
@@ -74,15 +75,16 @@ def main():
  # optionally resume from a checkpoint
  if args.resume:
  if os.path.isfile(args.resume):
- print("=> loading checkpoint '{}'".format(args.resume))
+ print("=> loading checkpoint '{}'".format(args.resume), flush=True)
  checkpoint = torch.load(args.resume)
  args.start_epoch = checkpoint['epoch']
  best_prec1 = checkpoint['best_prec1']
  model.load_state_dict(checkpoint['state_dict'])
  print("=> loaded checkpoint '{}' (epoch {})"
- .format(args.evaluate, checkpoint['epoch']))
+ .format(args.evaluate, checkpoint['epoch']), flush=True)
  else:
- print("=> no checkpoint found at '{}'".format(args.resume))
+ print("=> no checkpoint found at '{}'".format(args.resume),
+ flush=True)
 
  cudnn.benchmark = True
 
@@ -189,7 +191,8 @@ def train(train_loader, model, criterion, optimizer, epoch):
  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
  epoch, i, len(train_loader), batch_time=batch_time,
- data_time=data_time, loss=losses, top1=top1, top5=top5))
+ data_time=data_time, loss=losses, top1=top1, top5=top5),
+ flush=True)
 
 
 def validate(val_loader, model, criterion):
@@ -228,10 +231,10 @@ def validate(val_loader, model, criterion):
  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
  i, len(val_loader), batch_time=batch_time, loss=losses,
- top1=top1, top5=top5))
+ top1=top1, top5=top5), flush=True)
 
  print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
- .format(top1=top1, top5=top5))
+ .format(top1=top1, top5=top5), flush=True)
 
  return top1.avg