small adjustments

pix2tex/model/settings/config-vit.yaml

@@ -29,7 +29,7 @@ max_seq_len: 512
 max_width: 672
 min_height: 32
 min_width: 32
-micro_batchsize: 64
+micro_batchsize: -1
 model_path: checkpoints_add
 name: pix2tex-vit
 num_layers: 4

pix2tex/model/settings/config.yaml

@@ -6,7 +6,7 @@ backbone_layers:
 betas:
 - 0.9
 - 0.999
-batchsize: 10
+batchsize: 64
 bos_token: 1
 channels: 1
 data: dataset/data/train.pkl

pix2tex/train.py

@@ -12,23 +12,7 @@
 from pix2tex.eval import evaluate
 from pix2tex.models import get_model
 # from pix2tex.utils import *
-from pix2tex.utils import in_model_path, parse_args, seed_everything, get_optimizer, get_scheduler
-
-
-def gpu_memory_check(model, args):
- # check if largest batch can be handled by system
- try:
- batchsize = args.batchsize if args.get('micro_batchsize', -1) == -1 else args.micro_batchsize
- for _ in range(5):
- im = torch.empty(batchsize, args.channels, args.max_height, args.min_height, device=args.device).float()
- seq = torch.randint(0, args.num_tokens, (batchsize, args.max_seq_len), device=args.device).long()
- loss = model.data_parallel(im, device_ids=args.gpu_devices, tgt_seq=seq)
- loss.sum().backward()
- except RuntimeError:
- raise RuntimeError("The system cannot handle a batch size of %i for the maximum image size (%i, %i). Try to use a smaller micro batchsize." % (batchsize, args.max_height, args.max_width))
- model.zero_grad()
- with torch.cuda.device(args.device):torch.cuda.empty_cache()
- del im, seq
+from pix2tex.utils import in_model_path, parse_args, seed_everything, get_optimizer, get_scheduler, gpu_memory_check
 
 
 def train(args):
@@ -40,13 +24,15 @@ def train(args):
  valdataloader.update(**valargs)
  device = args.device
  model = get_model(args)
- gpu_memory_check(model, args)
- if args.load_chkpt is not None:
- model.load_state_dict(torch.load(args.load_chkpt, map_location=device))
+ if torch.cuda.is_available() and not args.no_cuda:
+ gpu_memory_check(model, args)
  max_bleu, max_token_acc = 0, 0
  out_path = os.path.join(args.model_path, args.name)
  os.makedirs(out_path, exist_ok=True)
 
+ if args.load_chkpt is not None:
+ model.load_state_dict(torch.load(args.load_chkpt, map_location=device))
+
  def save_models(e, step=0):
  torch.save(model.state_dict(), os.path.join(out_path, '%s_e%02d_step%02d.pth' % (args.name, e+1, step)))
  yaml.dump(dict(args), open(os.path.join(out_path, 'config.yaml'), 'w+'))
@@ -88,9 +74,9 @@ def save_models(e, step=0):
  wandb.log({'train/epoch': e+1})
  except KeyboardInterrupt:
  if e >= 2:
- save_models(e)
+ save_models(e, step=i)
  raise KeyboardInterrupt
- save_models(e)
+ save_models(e, step=len(dataloader))
 
 
 if __name__ == '__main__':

pix2tex/utils/utils.py

@@ -52,26 +52,44 @@ def seed_everything(seed: int):
 def parse_args(args, **kwargs) -> Munch:
  args = Munch({'epoch': 0}, **args)
  kwargs = Munch({'no_cuda': False, 'debug': False}, **kwargs)
+ args.update(kwargs)
  args.wandb = not kwargs.debug and not args.debug
- args.device = get_device(args, kwargs)
+ args.device = get_device(args, kwargs.no_cuda)
  args.max_dimensions = [args.max_width, args.max_height]
  args.min_dimensions = [args.get('min_width', 32), args.get('min_height', 32)]
  if 'decoder_args' not in args or args.decoder_args is None:
  args.decoder_args = {}
  return args
 
 
-def get_device(args, kwargs):
+def get_device(args, no_cuda=False):
  device = 'cpu'
  available_gpus = torch.cuda.device_count()
- args.gpu_devices = args.gpu_devices if args.get('gpu_devices', False) else range(available_gpus)
- if available_gpus > 0 and not kwargs.no_cuda:
+ args.gpu_devices = args.gpu_devices if args.get('gpu_devices', False) else list(range(available_gpus))
+ if available_gpus > 0 and not no_cuda:
  device = 'cuda:%d' % args.gpu_devices[0] if args.gpu_devices else 0
  assert available_gpus >= len(args.gpu_devices), "Available %d gpu, but specified gpu %s." % (available_gpus, ','.join(map(str, args.gpu_devices)))
- assert max(args.gpu_devices) < available_gpus, "legal gpu_devices should in [%s], received [%s]" % (','.join(map(str, range(available_gpus))),','.join(map(str, args.gpu_devices)))
+ assert max(args.gpu_devices) < available_gpus, "legal gpu_devices should in [%s], received [%s]" % (','.join(map(str, range(available_gpus))), ','.join(map(str, args.gpu_devices)))
  return device
 
 
+def gpu_memory_check(model, args):
+ # check if largest batch can be handled by system
+ try:
+ batchsize = args.batchsize if args.get('micro_batchsize', -1) == -1 else args.micro_batchsize
+ for _ in range(5):
+ im = torch.empty(batchsize, args.channels, args.max_height, args.min_height, device=args.device).float()
+ seq = torch.randint(0, args.num_tokens, (batchsize, args.max_seq_len), device=args.device).long()
+ loss = model.data_parallel(im, device_ids=args.gpu_devices, tgt_seq=seq)
+ loss.sum().backward()
+ except RuntimeError:
+ raise RuntimeError("The system cannot handle a batch size of %i for the maximum image size (%i, %i). Try to use a smaller micro batchsize." % (batchsize, args.max_height, args.max_width))
+ model.zero_grad()
+ with torch.cuda.device(args.device):
+ torch.cuda.empty_cache()
+ del im, seq
+
+
 def token2str(tokens, tokenizer) -> list:
  if len(tokens.shape) == 1:
  tokens = tokens[None, :]