def main(config):
from mxnet import nd
from mxnet.gluon.loss import CTCLoss
from models import get_model
from data_loader import get_dataloader
from trainer import Trainer
from utils import get_ctx, load
if os.path.isfile(config['dataset']['alphabet']):
config['dataset']['alphabet'] = ''.join(
load(config['dataset']['alphabet']))
prediction_type = config['arch']['args']['prediction']['type']
num_class = len(config['dataset']['alphabet'])
# loss 设置
if prediction_type == 'CTC':
criterion = CTCLoss()
else:
raise NotImplementedError
ctx = get_ctx(config['trainer']['gpus'])
model = get_model(num_class, ctx, config['arch']['args'])
model.hybridize()
model.initialize(ctx=ctx)
img_h, img_w = 32, 100
for process in config['dataset']['train']['dataset']['args'][
'pre_processes']:
if process['type'] == "Resize":
img_h = process['args']['img_h']
img_w = process['args']['img_w']
break
img_channel = 3 if config['dataset']['train']['dataset']['args'][
'img_mode'] != 'GRAY' else 1
sample_input = nd.zeros((2, img_channel, img_h, img_w), ctx[0])
num_label = model.get_batch_max_length(sample_input)
train_loader = get_dataloader(config['dataset']['train'], num_label,
config['dataset']['alphabet'])
assert train_loader is not None
if 'validate' in config['dataset']:
validate_loader = get_dataloader(config['dataset']['validate'],
num_label,
config['dataset']['alphabet'])
else:
validate_loader = None
config['lr_scheduler']['args']['step'] *= len(train_loader)
trainer = Trainer(config=config,
model=model,
criterion=criterion,
train_loader=train_loader,
validate_loader=validate_loader,
sample_input=sample_input,
ctx=ctx)
trainer.train()
def main(config):
if os.path.isfile(config['data_loader']['args']['dataset']['alphabet']):
config['data_loader']['args']['dataset']['alphabet'] = str(
np.load(config['data_loader']['args']['dataset']['alphabet']))
prediction_type = config['arch']['args']['prediction']['type']
num_class = len(config['data_loader']['args']['dataset']['alphabet'])
# loss 设置
if prediction_type == 'CTC':
criterion = CTCLoss()
else:
raise NotImplementedError
ctx = get_ctx(config['trainer']['gpus'])
model = get_model(num_class, config['arch']['args'])
model.hybridize()
model.initialize(ctx=ctx)
img_w = config['data_loader']['args']['dataset']['img_w']
img_h = config['data_loader']['args']['dataset']['img_h']
train_loader, val_loader = get_dataloader(
config['data_loader']['type'],
config['data_loader']['args'],
num_label=model.get_batch_max_length(img_h=img_h, img_w=img_w,
ctx=ctx))
config['lr_scheduler']['args']['step'] *= len(train_loader)
trainer = Trainer(config=config,
model=model,
criterion=criterion,
train_loader=train_loader,
val_loader=val_loader,
ctx=ctx)
trainer.train()
def forward(self, x):
out = x
for i, b in enumerate(self.net):
out = b(out)
if self.verbose:
print('Block %d output: %s' % (i + 1, out.shape))
return out
net = GoogleNet(10, verbose=True)
net.initialize()
x = nd.random.uniform(shape=(4, 3, 96, 96))
y = net(x)
ctx = get_ctx()
net = GoogleNet(10, verbose=False)
net.initialize(ctx=ctx, init=init.Xavier())
batch_size = 64
train_data, test_data = load_data_fashion_mnist(batch_size=batch_size,
transform=transform_resize)
softmax_xentropy = gluon.loss.SoftmaxCrossEntropyLoss()
trainer = gluon.Trainer(params=net.collect_params(),
optimizer='sgd',
optimizer_params={'learning_rate': 0.01})
train(train_data,
test_data,
net,
softmax_xentropy,