def __init__(self, config, pretrain=True):
self.config = config
self.model, self.vocab = build_model(config)
self.device = config['device']
self.num_epochs = config['trainer']['epochs']
self.data_root = config['trainer']['data_root']
self.train_annotation = config['trainer']['train_annotation']
self.valid_annotation = config['trainer']['valid_annotation']
self.batch_size = config['trainer']['batch_size']
self.print_every = config['trainer']['print_every']
self.valid_every = config['trainer']['valid_every']
self.checkpoint = config['trainer']['checkpoint']
self.export_weights = config['trainer']['export']
self.metrics = config['trainer']['metrics']
if pretrain:
download_weights(**config['pretrain'], quiet=config['quiet'])
self.model.load_state_dict(torch.load(config['pretrain']['cached'], map_location=torch.device(self.device)))
self.epoch = 0
self.iter = 0
self.optimizer = ScheduledOptim(
Adam(self.model.parameters(), betas=(0.9, 0.98), eps=1e-09),
0.2, config['transformer']['d_model'], config['optimizer']['n_warmup_steps'])
self.criterion = nn.CrossEntropyLoss(ignore_index=0)
self.train_gen = DataGen(self.data_root, self.train_annotation, self.vocab, self.device)
if self.valid_annotation:
self.valid_gen = DataGen(self.data_root, self.valid_annotation, self.vocab, self.device)
self.train_losses = []
def __init__(self, config, pretrained=True):
self.config = config
self.model, self.vocab = build_model(config)
self.device = config['device']
self.num_iters = config['trainer']['iters']
self.beamsearch = config['predictor']['beamsearch']
self.data_root = config['dataset']['data_root']
self.train_annotation = config['dataset']['train_annotation']
self.valid_annotation = config['dataset']['valid_annotation']
self.dataset_name = config['dataset']['name']
self.batch_size = config['trainer']['batch_size']
self.print_every = config['trainer']['print_every']
self.valid_every = config['trainer']['valid_every']
self.checkpoint = config['trainer']['checkpoint']
self.export_weights = config['trainer']['export']
self.metrics = config['trainer']['metrics']
logger = config['trainer']['log']
if logger:
self.logger = Logger(logger)
if pretrained:
weight_file = download_weights(**config['pretrain'],
quiet=config['quiet'])
self.load_weights(weight_file)
self.iter = 0
self.optimizer = AdamW(self.model.parameters(),
betas=(0.9, 0.98),
eps=1e-09)
self.scheduler = OneCycleLR(self.optimizer, **config['optimizer'])
# self.optimizer = ScheduledOptim(
# Adam(self.model.parameters(), betas=(0.9, 0.98), eps=1e-09),
# #config['transformer']['d_model'],
# 512,
# **config['optimizer'])
self.criterion = LabelSmoothingLoss(len(self.vocab),
padding_idx=self.vocab.pad,
smoothing=0.1)
transforms = ImgAugTransform()
self.train_gen = self.data_gen('train_{}'.format(self.dataset_name),
self.data_root,
self.train_annotation,
transform=transforms)
if self.valid_annotation:
self.valid_gen = self.data_gen(
'valid_{}'.format(self.dataset_name), self.data_root,
self.valid_annotation)
self.train_losses = []
def __init__(self, config):
device = config['device']
model, vocab = build_model(config)
weights = './models/reader/transformerocr.pth'
if config['weights'].startswith('http'):
weights = download_weights(config['weights'])
else:
weights = config['weights']
model.load_state_dict(
torch.load(weights, map_location=torch.device(device)))
self.config = config
self.model = model
self.vocab = vocab
def __init__(self, config, quanti=False):
device = config['device']
model, vocab = build_model(config)
weights = '/tmp/weights.pth'
if config['weights'].startswith('http'):
weights = download_weights(config['weights'])
else:
weights = config['weights']
model.load_state_dict(
torch.load(weights, map_location=torch.device(device)))
self.config = config
self.model = model
self.vocab = vocab
def __init__(self, config):
device = config['device']
model, vocab = build_model(config)
if config['weights'].startswith('http'):
weights = download_weights(config['weights'])
else:
weights = config['weights']
try:
model.load_state_dict(
torch.load(weights,
map_location=torch.device(device))['state_dict'])
except:
model.load_state_dict(
torch.load(weights, map_location=torch.device(device)))
self.config = config
self.model = model
self.vocab = vocab
def __init__(self, config, pretrained=True):
self.config = config
self.model, self.vocab = build_model(config)
self.device = config['device']
self.num_iters = config['trainer']['iters']
self.beamsearch = config['predictor']['beamsearch']
self.data_root = config['dataset']['data_root']
self.train_annotation = config['dataset']['train_annotation']
self.valid_annotation = config['dataset']['valid_annotation']
self.dataset_name = config['dataset']['name']
self.batch_size = config['trainer']['batch_size']
self.print_every = config['trainer']['print_every']
self.valid_every = config['trainer']['valid_every']
self.checkpoint = config['trainer']['checkpoint']
self.export_weights = config['trainer']['export']
self.metrics = config['trainer']['metrics']
logger = config['trainer']['log']
if logger:
self.logger = Logger(logger)
if pretrained:
download_weights(**config['pretrain'], quiet=config['quiet'])
state_dict = torch.load(config['pretrain']['cached'],
map_location=torch.device(self.device))
for name, param in self.model.named_parameters():
if state_dict[name].shape != param.shape:
print('{} missmatching shape'.format(name))
del state_dict[name]
self.model.load_state_dict(state_dict, strict=False)
self.iter = 0
self.optimizer = ScheduledOptim(
Adam(self.model.parameters(), betas=(0.9, 0.98), eps=1e-09),
config['transformer']['d_model'], **config['optimizer'])
# self.criterion = nn.CrossEntropyLoss(ignore_index=0)
self.criterion = LabelSmoothingLoss(len(self.vocab),
padding_idx=self.vocab.pad,
smoothing=0.1)
transforms = torchvision.transforms.Compose([
torchvision.transforms.ColorJitter(brightness=.1,
contrast=.1,
hue=.1,
saturation=.1),
torchvision.transforms.RandomAffine(degrees=0,
scale=(3 / 4, 4 / 3))
])
self.train_gen = self.data_gen('train_{}'.format(self.dataset_name),
self.data_root,
self.train_annotation,
transform=transforms)
if self.valid_annotation:
self.valid_gen = self.data_gen(
'valid_{}'.format(self.dataset_name), self.data_root,
self.valid_annotation)
self.train_losses = []