def main():
start_time = time()
net = Net().to(setting.device)
net.train()
print('The modle has been initialized.')
# define loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
train_dataloader = data_preprocess.get_train_dataloader()
for epoch in tqdm(range(num_epochs)):
start = time()
# for train accuracy
# total = setting.train_img_nums
# correct = 0
for batch_idx, (imgs, labels) in enumerate(train_dataloader):
# imgs = Variable(imgs)
# label = Variable(labels)
imgs, labels = imgs.to(setting.device), labels.to(setting.device)
labels = labels.long()
labels_ohe_predict = net(imgs)
loss = 0
for i in range(setting.char_num):
one_label = labels[:, i * setting.pool_length:(i + 1) *
setting.pool_length]
one_class = one_label.argmax(dim=1)
one_predict_label = labels_ohe_predict[:,
i * setting.pool_length:
(i + 1) *
setting.pool_length]
one_loss = criterion(one_predict_label, one_class)
loss += one_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
# for single in range(labels_ohe_predict.shape[0]):
# single_labels_ohe_predict = labels_ohe_predict[single, :]
# predict_label = ''
# # get predict_label
# for slice in range(setting.char_num):
# char = ohe.num2char[np.argmax(
# single_labels_ohe_predict[slice*setting.pool_length:(slice+1)*setting.pool_length].cpu().data.numpy())]
# predict_label += char
# # get true label
# true_label = ohe.decode(labels[single, :].cpu().numpy())
# if predict_label == true_label:
# correct += 1
end = time()
print('epoch: {}, time: {:.2f}s loss: {:.04}'.format(
epoch, end - start, loss.item()))
# print('epoch: {}, time: {:.2f}s loss: {:.04} accuracy: {}/{} -- {:.4f}'.format(
# epoch, end-start, loss.item(), correct, total, correct/total))
torch.save(net.state_dict(), './model.pt')
finnal_time = time()
print('End at {}, cost {:.0f}s'.format(finnal_time,
finnal_time - start_time))
def train(loaders, save_path):
"""returns trained model"""
# Initialize custom defined cnn
model = Net()
use_cuda = torch.cuda.is_available()
if use_cuda:
model.cuda()
# cross entropy loss for classification task
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=config.lr)
# initialize tracker for minimum validation loss
valid_loss_min = np.Inf
n_epochs = config.n_epochs
for epoch in range(1, n_epochs + 1):
# initialize variables to monitor training and validation loss
train_loss = 0.0
valid_loss = 0.0
model.train()
for batch_idx, (data, target) in enumerate(loaders['train']):
# move to GPU
if use_cuda:
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
# average training loss
train_loss += (1 / (batch_idx + 1)) * (loss.data - train_loss)
# vaidation
model.eval()
for batch_idx, (data, target) in enumerate(loaders['valid']):
# move to GPU
if use_cuda:
data, target = data.cuda(), target.cuda()
## update the average validation loss
output = model(data)
loss = criterion(output, target)
valid_loss += (1 / (batch_idx + 1)) * (loss.data - valid_loss)
# print training/validation statistics
print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.
format(epoch, train_loss, valid_loss))
# save the model if validation loss has decreased
if valid_loss <= valid_loss_min:
torch.save(model.state_dict(), save_path)
# Updating the validation loss minimum
valid_loss_min = valid_loss
# return trained model
return model
device = torch.device("cuda" if use_cuda else "cpu")
print(device)
net = Net().to(device)
print(net)
optimizer = optim.Adam(net.parameters())
epochs = 50
# net.load_state_dict(torch.load('model.net'))
for epoch in range(1, epochs + 1):
train(net, device, train_loader, optimizer, epoch)
test(net, device, validation_loader)
torch.save(net.state_dict(), "model_2.net")
for i_batch, sample_batched in enumerate(validation_loader):
net.eval()
data, target = sample_batched['image'].float().to(
device), sample_batched['keypoints'].float().to(device)
out = net(data)
if i_batch == 1:
for i in range(5):
plt.figure()
plt.imshow(data[i].cpu().numpy().transpose(1, 2, 0))
target_i = target[i].cpu().numpy().reshape(-1, 2)
out_i = out[i].cpu().detach().numpy().reshape(-1, 2)
plt.scatter(out_i[:, 0], out_i[:, 1], s=10, marker='.', c='r')
# plt.scatter(target_i[:, 0], target_i[:, 1], s=10, marker='.', c='b')
plt.axis('off')
shuffle=True,
drop_last = True)
model = Net()
if torch.cuda.is_available():
model.cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(),
lr = args.learning_rate)
for epoch in range(args.epoch):
start_time = time.time()
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
inputs, labels = data
if torch.cuda.is_available():
inputs = inputs.cuda()
labels = labels.cuda()
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print('Epoch no: %d, Epoch loss: %.3f, Epoch time = %.3f' %
(epoch + 1, running_loss, time.time() - start_time))
print('Finished Training')
torch.save(model.state_dict(),
args.cnn_model_location)