def train(args):
if not os.path.exists('checkpoints'):
os.mkdir('checkpoints')
# Setup Augmentations
data_aug = transforms.Compose([
transforms.ToPILImage(),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomAffine(degrees=10, translate=(0.05, 0.05), scale=(0.95, 1.05)),
])
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
t_loader = data_loader(data_path, is_transform=True, split='train', version='simplified', img_size=(args.img_rows, args.img_cols), augmentations=data_aug, train_fold_num=args.train_fold_num, num_train_folds=args.num_train_folds, seed=args.seed)
v_loader = data_loader(data_path, is_transform=True, split='val', version='simplified', img_size=(args.img_rows, args.img_cols), num_val=args.num_val, seed=args.seed)
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
n_classes = t_loader.n_classes
trainloader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=2, shuffle=True, pin_memory=True, drop_last=True)
valloader = data.DataLoader(v_loader, batch_size=args.batch_size, num_workers=2, pin_memory=True)
# Setup Metrics
running_metrics = runningScore(n_classes)
# Setup Model
# model = get_model(args.arch, n_classes, use_cbam=args.use_cbam)
model = torchvision.models.mobilenet_v2(pretrained=True)
num_ftrs = model.last_channel
model.classifier = nn.Sequential(
nn.Dropout(0.2),
nn.Linear(num_ftrs, n_classes),
)
model.cuda()
# Check if model has custom optimizer / loss
if hasattr(model, 'optimizer'):
optimizer = model.optimizer
else:
##optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=args.l_rate, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer = torch.optim.Adam(model.parameters(), lr=args.l_rate, weight_decay=args.weight_decay)
# if args.num_cycles > 0:
# len_trainloader = int(5e6) # 4960414
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.num_train_folds*len_trainloader//args.num_cycles, eta_min=args.l_rate*1e-1)
# else:
# scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[2, 4, 6, 8], gamma=0.5)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,mode='min',factor=0.5,patience=5,cooldown=5,min_lr=1e-7)
if hasattr(model, 'loss'):
print('Using custom loss')
loss_fn = model.loss
else:
loss_fn = F.cross_entropy
start_epoch = 0
if args.resume is not None:
if os.path.isfile(args.resume):
print("Loading model and optimizer from checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model_dict = model.state_dict()
if checkpoint.get('model_state', -1) == -1:
model_dict.update(convert_state_dict(checkpoint, load_classifier=args.load_classifier))
else:
model_dict.update(convert_state_dict(checkpoint['model_state'], load_classifier=args.load_classifier))
print("Loaded checkpoint '{}' (epoch {}, mapk {:.5f}, top1_acc {:7.3f}, top2_acc {:7.3f} top3_acc {:7.3f})"
.format(args.resume, checkpoint['epoch'], checkpoint['mapk'], checkpoint['top1_acc'], checkpoint['top2_acc'], checkpoint['top3_acc']))
model.load_state_dict(model_dict)
if checkpoint.get('optimizer_state', None) is not None:
optimizer.load_state_dict(checkpoint['optimizer_state'])
start_epoch = checkpoint['epoch']
else:
print("No checkpoint found at '{}'".format(args.resume))
loss_sum = 0.0
for epoch in range(start_epoch, args.n_epoch):
start_train_time = timeit.default_timer()
model.train()
optimizer.zero_grad()
for i, (images, labels, recognized, _) in enumerate(trainloader):
images = images.cuda()
labels = labels.cuda()
recognized = recognized.cuda()
outputs = model(images)
loss = (loss_fn(outputs, labels.view(-1), ignore_index=t_loader.ignore_index, reduction='none') * recognized.view(-1)).mean()
# loss = loss / float(args.iter_size) # Accumulated gradients
loss_sum = loss_sum + loss
loss.backward()
if (i+1) % args.print_train_freq == 0:
print("Epoch [%d/%d] Iter [%6d/%6d] Loss: %.4f" % (epoch+1, args.n_epoch, i+1, len(trainloader), loss_sum))
if (i+1) % args.iter_size == 0 or i == len(trainloader) - 1:
optimizer.step()
optimizer.zero_grad()
loss_sum = 0.0
mapk_val = AverageMeter()
top1_acc_val = AverageMeter()
top2_acc_val = AverageMeter()
top3_acc_val = AverageMeter()
mean_loss_val = AverageMeter()
model.eval()
with torch.no_grad():
for i_val, (images_val, labels_val, recognized_val, _) in tqdm(enumerate(valloader)):
images_val = images_val.cuda()
labels_val = labels_val.cuda()
recognized_val = recognized_val.cuda()
outputs_val = model(images_val)
loss_val = (loss_fn(outputs_val, labels_val.view(-1), ignore_index=v_loader.ignore_index, reduction='none') * recognized_val.view(-1)).mean()
mean_loss_val.update(loss_val, n=images_val.size(0))
_, pred = outputs_val.topk(k=3, dim=1, largest=True, sorted=True)
running_metrics.update(labels_val, pred[:, 0])
acc1, acc2, acc3 = accuracy(outputs_val, labels_val, topk=(1, 2, 3))
top1_acc_val.update(acc1, n=images_val.size(0))
top2_acc_val.update(acc2, n=images_val.size(0))
top3_acc_val.update(acc3, n=images_val.size(0))
mapk_v = mapk(labels_val, pred, k=3)
mapk_val.update(mapk_v, n=images_val.size(0))
print('Mean Average Precision (MAP) @ 3: {:.5f}'.format(mapk_val.avg))
print('Top 3 accuracy: {:7.3f} / {:7.3f} / {:7.3f}'.format(top1_acc_val.avg, top2_acc_val.avg, top3_acc_val.avg))
print('Mean val loss: {:.4f}'.format(mean_loss_val.avg))
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
print(k, v)
#for i in range(n_classes):
# print(i, class_iou[i])
scheduler.step(mean_loss_val.avg)
state = {'epoch': epoch+1,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(),
'mapk': mapk_val.avg,
'top1_acc': top1_acc_val.avg,
'top2_acc': top2_acc_val.avg,
'top3_acc': top3_acc_val.avg,}
torch.save(state, "checkpoints/{}_{}_{}_{}x{}_{}-{}-{}_model.pth".format(args.arch, args.dataset, epoch+1, args.img_rows, args.img_cols, args.train_fold_num, args.num_train_folds, args.num_val))
running_metrics.reset()
mapk_val.reset()
top1_acc_val.reset()
top2_acc_val.reset()
top3_acc_val.reset()
mean_loss_val.reset()
elapsed_train_time = timeit.default_timer() - start_train_time
print('Training time (epoch {0:5d}): {1:10.5f} seconds'.format(epoch+1, elapsed_train_time))
def test(args):
model_file_name = os.path.split(args.model_path)[1]
model_name = model_file_name[:model_file_name.find('_')]
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
loader = data_loader(data_path,
split=args.split,
is_transform=True,
img_size=(args.img_rows, args.img_cols),
no_gt=args.no_gt,
seed=args.seed)
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
n_classes = loader.n_classes
testloader = data.DataLoader(loader,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True)
# Setup Model
model = torchvision.models.mobilenet_v2(pretrained=True)
num_ftrs = model.last_channel
model.classifier = nn.Sequential(
nn.Dropout(0.2),
nn.Linear(num_ftrs, n_classes),
)
model.cuda()
checkpoint = torch.load(args.model_path)
state = convert_state_dict(checkpoint['model_state'])
model_dict = model.state_dict()
model_dict.update(state)
model.load_state_dict(model_dict)
print(
"Loaded checkpoint '{}' (epoch {}, mapk {:.5f}, top1_acc {:7.3f}, top2_acc {:7.3f} top3_acc {:7.3f})"
.format(args.model_path, checkpoint['epoch'], checkpoint['mapk'],
checkpoint['top1_acc'], checkpoint['top2_acc'],
checkpoint['top3_acc']))
running_metrics = runningScore(n_classes)
pred_dict = collections.OrderedDict()
mapk = AverageMeter()
model.eval()
with torch.no_grad():
for i, (images, labels, _, names) in tqdm(enumerate(testloader)):
plt.imshow((images[0].numpy().transpose(1, 2, 0) -
np.min(images[0].numpy().transpose(1, 2, 0))) /
(np.max(images[0].numpy().transpose(1, 2, 0) -
np.min(images[0].numpy().transpose(1, 2, 0)))))
plt.show()
images = images.cuda()
if args.tta:
images_flip = flip(images, dim=3)
outputs = model(images)
if args.tta:
outputs_flip = model(images_flip)
prob = F.softmax(outputs, dim=1)
if args.tta:
prob_flip = F.softmax(outputs_flip, dim=1)
prob = (prob + prob_flip) / 2.0
_, pred = prob.topk(k=3, dim=1, largest=True, sorted=True)
for k in range(images.size(0)):
pred_dict[int(names[0][k])] = loader.encode_pred_name(
pred[k, :])
if not args.no_gt:
running_metrics.update(labels, pred)
mapk_val = mapk(labels, pred, k=3)
mapk.update(mapk_val, n=images.size(0))
print('Mean Average Precision (MAP) @ 3: {:.5f}'.format(mapk.avg))
if not args.no_gt:
print('Mean Average Precision (MAP) @ 3: {:.5f}'.format(mapk.avg))
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
print(k, v)
#for i in range(n_classes):
# print(i, class_iou[i])
running_metrics.reset()
mapk.reset()
# Create submission
sub = pd.DataFrame.from_dict(pred_dict, orient='index')
sub.index.names = ['key_id']
sub.columns = ['word']
sub.to_csv('{}_{}x{}.csv'.format(args.split, args.img_rows, args.img_cols))