def main(args):
# 1. prepare data & models
train_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
TransformByKeys(transforms.ToPILImage(), ("image",)),
TransformByKeys(transforms.ToTensor(), ("image",)),
TransformByKeys(transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ("image",)),
])
print("Reading data...")
train_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'), train_transforms, split="train")
train_dataloader = data.DataLoader(train_dataset, batch_size=args.batch_size, num_workers=4, pin_memory=True,
shuffle=True, drop_last=True)
val_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'), train_transforms, split="val")
val_dataloader = data.DataLoader(val_dataset, batch_size=args.batch_size, num_workers=4, pin_memory=True,
shuffle=False, drop_last=False)
print("Creating model...")
device = torch.device("cuda: 0") if args.gpu else torch.device("cpu")
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 2 * NUM_PTS, bias=True)
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate, amsgrad=True)
loss_fn = fnn.mse_loss
# 2. train & validate
print("Ready for training...")
best_val_loss = np.inf
for epoch in range(args.epochs):
train_loss = train(model, train_dataloader, loss_fn, optimizer, device=device)
val_loss = validate(model, val_dataloader, loss_fn, device=device)
print("Epoch #{:2}:\ttrain loss: {:5.2}\tval loss: {:5.2}".format(epoch, train_loss, val_loss))
if val_loss < best_val_loss:
best_val_loss = val_loss
with open(f"{args.name}_best.pth", "wb") as fp:
torch.save(model.state_dict(), fp)
# 3. predict
test_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'test'), train_transforms, split="test")
test_dataloader = data.DataLoader(test_dataset, batch_size=args.batch_size, num_workers=4, pin_memory=True,
shuffle=False, drop_last=False)
with open(f"{args.name}_best.pth", "rb") as fp:
best_state_dict = torch.load(fp, map_location="cpu")
model.load_state_dict(best_state_dict)
test_predictions = predict(model, test_dataloader, device)
with open(f"{args.name}_test_predictions.pkl", "wb") as fp:
pickle.dump({"image_names": test_dataset.image_names,
"landmarks": test_predictions}, fp)
create_submission(args.data, test_predictions, f"{args.name}_submit.csv")
def main(args):
# 1. prepare data & models
train_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
AffineAugmenter(min_scale=0.9, max_offset=0.1, rotate=True),
BrightnessContrastAugmenter(brightness=0.3, contrast=0.3),
BlurAugmenter(max_kernel=5),
TransformByKeys(transforms.ToPILImage(), ("image", )),
TransformByKeys(transforms.ToTensor(), ("image", )),
TransformByKeys(
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
("image", )),
])
test_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
TransformByKeys(transforms.ToPILImage(), ("image", )),
TransformByKeys(transforms.ToTensor(), ("image", )),
TransformByKeys(
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
("image", )),
])
print("Reading data...")
train_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'),
train_transforms,
split="train")
train_dataloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=True,
drop_last=True)
val_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'),
test_transforms,
split="val")
val_dataloader = data.DataLoader(val_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=False,
drop_last=False)
print("Creating model...")
device = torch.device("cuda: 0") if args.gpu else torch.device("cpu")
model = models.resnext50_32x4d(pretrained=True)
# for param in model.parameters():
# param.requires_grad = False
model.fc = nn.Linear(model.fc.in_features, 2 * NUM_PTS, bias=True)
# model.fc = nn.Sequential(
# # nn.BatchNorm1d(model.fc.in_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
# # nn.Linear(model.fc.in_features, model.fc.in_features, bias=True),
# # nn.ReLU(),
# nn.BatchNorm1d(model.fc.in_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
# nn.Linear(model.fc.in_features, 2 * NUM_PTS, bias=True))
model.to(device)
# optimizer = optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=0.01, amsgrad=True)
optimizer = RAdam(model.parameters(),
lr=args.learning_rate) # , weight_decay=0.01)
optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=3)
# optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.2)
loss_fn = fnn.mse_loss
# 2. train & validate
print("Ready for training...")
best_val_loss = np.inf
for epoch in range(args.epochs):
train_loss = train(model,
train_dataloader,
loss_fn,
optimizer,
device=device)
val_loss = validate(model, val_dataloader, loss_fn, device=device)
print("Epoch #{:2}:\ttrain loss: {:5.2}\tval loss: {:5.2}".format(
epoch, train_loss, val_loss))
if val_loss < best_val_loss:
best_val_loss = val_loss
with open(f"{args.name}_best.pth", "wb") as fp:
torch.save(model.state_dict(), fp)
# with open(f"{args.name}_{epoch}_{train_loss:7.4}_{val_loss:7.4}.pth", "wb") as fp:
# torch.save(model.state_dict(), fp)
# 3. predict
test_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'test'),
test_transforms,
split="test")
test_dataloader = data.DataLoader(test_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=False,
drop_last=False)
with open(f"{args.name}_best.pth", "rb") as fp:
best_state_dict = torch.load(fp, map_location="cpu")
model.load_state_dict(best_state_dict)
test_predictions = predict(model, test_dataloader, device)
with open(f"{args.name}_test_predictions.pkl", "wb") as fp:
pickle.dump(
{
"image_names": test_dataset.image_names,
"landmarks": test_predictions
}, fp)
create_submission(args.data, test_predictions, f"{args.name}_submit.csv")
if args.draw:
print("Drawing landmarks...")
directory = os.path.join("result",
test_dataset.image_names[0].split('.')[0])
if not os.path.exists(directory):
os.makedirs(directory)
random_idxs = np.random.choice(len(test_dataset.image_names),
size=1000,
replace=False)
for i, idx in enumerate(random_idxs, 1):
image = cv2.imread(test_dataset.image_names[idx])
image = draw_landmarks(image, test_predictions[idx])
cv2.imwrite(os.path.join("result", test_dataset.image_names[idx]),
image)
def main(args):
# 1. prepare data & models
# применение новых трансформаций не дало улучшения результатов
# единственное изменение это параметры нормализации
train_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
TransformByKeys(transforms.ToPILImage(), ("image", )),
TransformByKeys(transforms.ToTensor(), ("image", )),
TransformByKeys(
transforms.Normalize(
mean=[
0.485, 0.456, 0.406
], # средние значения и дисперсии взяты из документации pytorch,
std=[0.229, 0.224, 0.225]
), # с такими же значениями обучалась сеть на корпусе imagenet
("image", ),
),
])
device = torch.device("cuda: 0") if args.gpu else torch.device("cpu")
print("Creating model...")
model = models.resnext50_32x4d(pretrained=True)
in_features = model.fc.in_features
fc = nn.Sequential(nn.Linear(in_features, 2 * NUM_PTS), ) # новая "голова"
model.fc = fc
state_dict = None
# если есть сеть дообученная на датасете из контеста
if args.pretrained_model:
print(f"Load best_state_dict {args.pretrained_model}")
state_dict = torch.load(args.pretrained_model)
model.load_state_dict(state_dict)
del state_dict
model.to(device)
print(model)
factor = 0.1**(1 / 2) # уменьшающий фактор для lr
# оптимизатора выбран AdamW, с небольшой нормализацией весов
optimizer = optim.AdamW(model.parameters(),
lr=args.learning_rate,
amsgrad=True,
weight_decay=0.05)
loss_fn = fnn.mse_loss
# изменения lr происходит при помощи ReduceLROnPlateau
scheduler = ReduceLROnPlateau(
optimizer,
mode='min',
patience=1,
factor=factor,
)
print(loss_fn)
print(optimizer)
print(scheduler)
print("Reading data...")
print("Read train landmark dataset")
train_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'),
train_transforms,
split="train")
print("Create picture loader for test dataset")
train_dataloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=0,
pin_memory=True,
shuffle=True,
drop_last=True)
print("Read val landmark dataset")
val_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'),
train_transforms,
split="val")
print("Create picture loader for val dataset")
val_dataloader = data.DataLoader(val_dataset,
batch_size=args.batch_size,
num_workers=0,
pin_memory=True,
shuffle=False,
drop_last=False)
# 2. train & validate
print("Ready for training...")
best_val_loss = np.inf
for epoch in range(args.epochs):
train_loss = train(model,
train_dataloader,
loss_fn,
optimizer,
device=device)
val_loss = validate(model, val_dataloader, loss_fn, device=device)
print("Epoch #{:2}:\ttrain loss: {:5.5}\tval loss: {:5.5}".format(
epoch + 1, train_loss, val_loss))
scheduler.step(val_loss)
if val_loss < best_val_loss:
best_val_loss = val_loss
with open(f"{args.name}_best.pth", "wb") as fp:
torch.save(model.state_dict(), fp)
# 3. predict
print("Start predict")
test_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'test'),
train_transforms,
split="test")
test_dataloader = data.DataLoader(test_dataset,
batch_size=args.batch_size,
num_workers=0,
pin_memory=True,
shuffle=False,
drop_last=False)
with open(f"{args.name}_best.pth", "rb") as fp:
best_state_dict = torch.load(fp, map_location="cpu")
model.load_state_dict(best_state_dict)
test_predictions = predict(model, test_dataloader, device)
with open(f"{args.name}_test_predictions.pkl", "wb") as fp:
pickle.dump(
{
"image_names": test_dataset.image_names,
"landmarks": test_predictions
}, fp)
create_submission(args.data, test_predictions, f"{args.name}_submit.csv")
def main(args):
print(torch.cuda.device_count(), 'gpus available')
# 0. Initializing training
if args.fold is not None:
if args.fold_prefix is None:
print('Please add fold-prefix to arguments')
return
folder_name = f'{args.name}_{args.fold_prefix // 10}{args.fold_prefix % 10}_fold{args.fold}'
checkpoint_path = os.path.join(args.data, 'checkpoints', folder_name)
log_path = os.path.join(args.data, 'logs', folder_name)
if not os.path.exists(checkpoint_path):
os.mkdir(checkpoint_path)
else:
for i in range(100):
folder_name = f'{args.name}_{i // 10}{i % 10}'
checkpoint_path = os.path.join(args.data, 'checkpoints', folder_name)
log_path = os.path.join(args.data, 'logs', folder_name)
if not os.path.exists(checkpoint_path):
os.mkdir(checkpoint_path)
break
if args.checkpoint is None:
training_state = {
'best_checkpoints': [],
'best_scores': [],
'epoch': []
}
else:
# loading checkpoint
from_checkpoint = f'{args.name}_{args.checkpoint // 10}{args.checkpoint % 10}'
parent_checkpoint_path = os.path.join(args.data, 'checkpoints', from_checkpoint)
training_state = torch.load(os.path.join(parent_checkpoint_path, 'training_state.pth'))
training_state['from_checkpoint'] = from_checkpoint
print(f'Using checkpoint {from_checkpoint}')
print(f'Results can be found in {folder_name}')
writer = SummaryWriter(log_dir=log_path)
# 1. prepare data & models
if args.name == 'senet154':
crop_size = 224
else:
crop_size = CROP_SIZE
train_transforms = transforms.Compose([
# HorizontalFlip(p=0.5),
ScaleMinSideToSize((crop_size, crop_size)),
CropCenter(crop_size),
TransformByKeys(transforms.ToPILImage(), ('image',)),
TransformByKeys(transforms.ToTensor(), ('image',)),
TransformByKeys(transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ('image',)),
])
test_transforms = transforms.Compose([
ScaleMinSideToSize((crop_size, crop_size)),
CropCenter(crop_size),
TransformByKeys(transforms.ToPILImage(), ('image',)),
TransformByKeys(transforms.ToTensor(), ('image',)),
TransformByKeys(transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ('image',)),
])
albu_transforms = albu.Compose([
albu.Blur(p=0.1),
albu.MultiplicativeNoise(p=0.1, per_channel=True),
albu.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=20, val_shift_limit=20, p=0.2),
albu.ChannelShuffle(p=0.2)
],
keypoint_params=albu.KeypointParams(format='xy'))
print('\nTransforms:')
print(albu_transforms)
print(train_transforms)
print('\nReading data...')
datasets = torch.load(os.path.join(args.data, 'datasets.pth'))
for d in datasets:
datasets[d].transforms = train_transforms
if args.fold is None:
print('Using predefined data split')
train_dataset = datasets['train_dataset']
val_dataset = datasets['val_dataset']
else:
print(f'Using fold {args.fold}')
train_dataset = FoldDatasetDataset(datasets['train_dataset'], datasets['val_dataset'], train_transforms,
albu_transforms, split='train', fold=args.fold, seed=42)
val_dataset = FoldDatasetDataset(datasets['train_dataset'], datasets['val_dataset'], train_transforms,
None, split='val', fold=args.fold, seed=42)
test_dataset = datasets['test_dataset']
test_dataset.transforms = test_transforms
train_dataloader = data.DataLoader(train_dataset, batch_size=args.batch_size, num_workers=32, pin_memory=True,
shuffle=True, drop_last=True)
val_dataloader = data.DataLoader(val_dataset, batch_size=args.batch_size, num_workers=32, pin_memory=True,
shuffle=False, drop_last=False)
print('Creating model...')
device = torch.device('cuda: 0') if args.gpu else torch.device('cpu')
if args.name == 'senet154':
model = pretrainedmodels.senet154(num_classes=1000, pretrained='imagenet')
model.last_linear = nn.Linear(model.last_linear.in_features, 2 * NUM_PTS, bias=True)
else:
model = models.resnext50_32x4d(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 2 * NUM_PTS, bias=True)
model = nn.DataParallel(model)
print(f'Using {torch.cuda.device_count()} gpus')
if args.checkpoint is not None:
model.load_state_dict(training_state['best_checkpoints'][0])
model.to(device)
# optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, nesterov=True)
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
# optimizer = RAdam(model.parameters(), lr=args.learning_rate)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
# scheduler = None
print(f'Optimizer: {optimizer}')
print(f'Scheduler: {scheduler}')
loss_fn = fnn.mse_loss
# 2. train & validate
print('Ready for training...')
if args.checkpoint is None:
start_epoch = 0
best_val_loss = np.inf
else:
start_epoch = training_state['epoch'][0]
best_val_loss = training_state['best_scores'][0]
for epoch in range(start_epoch, start_epoch + args.epochs):
train_loss = train(model, train_dataloader, loss_fn, optimizer, device, writer, epoch)
val_loss = validate(model, val_dataloader, loss_fn, device, writer, epoch, scheduler)
print('Epoch #{:2}:\ttrain loss: {:5.2}\tval loss: {:5.2}'.format(epoch, train_loss, val_loss))
print(f'Learning rate = {optimizer.param_groups[0]["lr"]}')
if len(training_state['best_scores']) == 0:
training_state['best_checkpoints'].append(model.state_dict())
training_state['best_scores'].append(val_loss)
training_state['epoch'].append(epoch)
with open(os.path.join(checkpoint_path, 'training_state.pth'), 'wb') as fp:
torch.save(training_state, fp)
elif len(training_state['best_scores']) < 3 or val_loss < training_state['best_scores'][-1]:
cur_val_index = 0
for cur_val_index in range(len(training_state['best_scores'])):
if val_loss < training_state['best_scores'][cur_val_index]:
break
training_state['best_scores'].insert(cur_val_index, val_loss)
training_state['best_checkpoints'].insert(cur_val_index, model.state_dict())
training_state['epoch'].insert(cur_val_index, epoch)
if len(training_state['best_scores']) > 3:
training_state['best_scores'] = training_state['best_scores'][:3]
training_state['best_checkpoints'] = training_state['best_checkpoints'][:3]
training_state['epoch'] = training_state['epoch'][:3]
with open(os.path.join(checkpoint_path, 'training_state.pth'), 'wb') as fp:
torch.save(training_state, fp)
if val_loss < best_val_loss:
best_val_loss = val_loss
with open(os.path.join(checkpoint_path, f'{args.name}_best.pth'), 'wb') as fp:
torch.save(model.state_dict(), fp)
print('Training finished')
print(f'Min val loss = {training_state["best_scores"]} at epoch {training_state["epoch"]}')
print()
# 3. predict
test_dataloader = data.DataLoader(test_dataset, batch_size=args.batch_size, num_workers=16, pin_memory=True,
shuffle=False, drop_last=False)
with open(os.path.join(checkpoint_path, f'{args.name}_best.pth'), 'rb') as fp:
best_state_dict = torch.load(fp, map_location='cpu')
model.load_state_dict(best_state_dict)
test_predictions = predict(model, test_dataloader, device)
with open(os.path.join(checkpoint_path, f'{args.name}_test_predictions.pkl'), 'wb') as fp:
pickle.dump({'image_names': test_dataset.image_names,
'landmarks': test_predictions}, fp)
create_submission(args.data, test_predictions, os.path.join(checkpoint_path, f'{args.name}_submit.csv'))
def main(args):
# 1. prepare data & models
train_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
TransformByKeys(transforms.ToPILImage(), ("image", )),
TransformByKeys(transforms.ToTensor(), ("image", )),
TransformByKeys(
transforms.Normalize(mean=[0.39963884, 0.31994772, 0.28253724],
std=[0.33419772, 0.2864468, 0.26987]),
("image", )),
])
print("Reading data...")
train_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'),
train_transforms,
split="train")
train_dataloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=True,
drop_last=True)
val_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'),
train_transforms,
split="val")
val_dataloader = data.DataLoader(val_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=False,
drop_last=False)
print("Creating model...")
device = torch.device("cuda: 0") if args.gpu else torch.device("cpu")
# model = models.wide_resnet101_2(pretrained=True)
# fc_layers = nn.Sequential(
# nn.Linear(model.fc.in_features, model.fc.in_features),
# nn.ReLU(inplace=True),
# nn.Dropout(p=0.1),
# nn.Linear(model.fc.in_features, 2 * NUM_PTS),
# nn.ReLU(inplace=True),
# nn.Dropout(p=0.1))
# model.fc = fc_layers
model = models.resnext101_32x8d(pretrained=True)
# Uncomment for learning with freezed feature extractor
# for param in model.parameters():
# param.requires_grad = False
model.fc = nn.Linear(model.fc.in_features, 2 * NUM_PTS, bias=True)
if args.checkpoint is not None:
model.load_state_dict(torch.load(args.checkpoint))
print('PRETRAIDED LOADED')
model.to(device)
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
amsgrad=True)
loss_fn = fnn.mse_loss
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
# 2. train & validate
writer = SummaryWriter()
print("Ready for training...")
best_val_loss = np.inf
for epoch in range(args.epochs):
train_loss = train(model,
train_dataloader,
loss_fn,
optimizer,
device=device,
epoch=epoch,
writer=writer)
val_loss = validate(model,
val_dataloader,
loss_fn,
device=device,
epoch=epoch,
writer=writer)
# if epoch > 0:
scheduler.step()
print(f"EPOCH {epoch} \n")
print("Epoch #{:2}:\ttrain loss: {:5.2}\tval loss: {:5.2}".format(
epoch, train_loss, val_loss))
if val_loss < best_val_loss:
best_val_loss = val_loss
with open(f"{args.name}_best.pth", "wb") as fp:
torch.save(model.state_dict(), fp)
# 3. predict
test_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'test'),
train_transforms,
split="test")
test_dataloader = data.DataLoader(test_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=False,
drop_last=False)
with open(f"{args.name}_best.pth", "rb") as fp:
best_state_dict = torch.load(fp, map_location="cpu")
model.load_state_dict(best_state_dict)
test_predictions = predict(model, test_dataloader, device)
with open(f"{args.name}_test_predictions.pkl", "wb") as fp:
pickle.dump(
{
"image_names": test_dataset.image_names,
"landmarks": test_predictions
}, fp)
create_submission(args.data, test_predictions, f"{args.name}_submit.csv")
def main(args):
# 1. prepare data & models
train_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
TransformByKeys(transforms.ToPILImage(), ("image", )),
TransformByKeys(transforms.ToTensor(), ("image", )),
TransformByKeys(
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]), ("image", )),
])
test_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
TransformByKeys(transforms.ToPILImage(), ("image", )),
TransformByKeys(transforms.ToTensor(), ("image", )),
TransformByKeys(
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]), ("image", )),
])
print(datetime.datetime.now())
print("Reading data...")
train_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'),
train_transforms,
split="train")
train_dataloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=True,
drop_last=True)
val_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'),
test_transforms,
split="val")
val_dataloader = data.DataLoader(val_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=False,
drop_last=False)
print(datetime.datetime.now())
print("Creating model...")
device = torch.device("cuda: 0") if args.gpu else torch.device("cpu")
model = torch.hub.load(
'facebookresearch/WSL-Images', 'resnext101_32x8d_wsl'
) # models.resnext50_32x4d(pretrained=True) # resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 2 * NUM_PTS, bias=True)
if os.path.isfile(f"{args.name}_best.pth"):
print("Loading saved model " + f"{args.name}_best.pth")
with open(f"{args.name}_best.pth", "rb") as fp:
best_state_dict = torch.load(fp, map_location="cpu")
model.load_state_dict(best_state_dict)
model.to(device)
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
amsgrad=True)
loss_fn = fnn.l1_loss # WingLoss() #fnn.mse_loss
loss_val = fnn.mse_loss
# 2. train & validate
print("Ready for training...")
best_val_loss = np.inf
for epoch in range(args.epochs):
train_loss = train(model,
train_dataloader,
loss_fn,
optimizer,
device=device)
val_loss = validate(model, val_dataloader, loss_fn, device=device)
val_loss_mse = validate(model, val_dataloader, loss_val, device=device)
print(
"Epoch #{:2}:\ttrain loss: {:5.4}\tval loss: {:5.4}\tval mse: {:5.4}"
.format(epoch, train_loss, val_loss, val_loss_mse))
if 1 == 1: #val_loss < best_val_loss: #save results of all epoch to check several at kaggle
best_val_loss = val_loss
with open(f"{args.name}_best.pth", "wb") as fp:
torch.save(model.state_dict(), fp)
with open(f"{args.name}_" + str(epoch) + ".pth", "wb") as fp:
torch.save(model.state_dict(), fp)
# 3. predict
print('Predict')
test_dataset = ThousandLandmarksDataset(os.path.join(
args.data, 'test'),
test_transforms,
split="test")
test_dataloader = data.DataLoader(test_dataset,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=False,
drop_last=False)
#with open(f"{args.name}_best.pth", "rb") as fp:
# best_state_dict = torch.load(fp, map_location="cpu")
# model.load_state_dict(best_state_dict)
test_predictions = predict(model, test_dataloader, device)
with open(f"{args.name}_test_predictions.pkl", "wb") as fp:
pickle.dump(
{
"image_names": test_dataset.image_names,
"landmarks": test_predictions
}, fp)
create_submission(args.data, test_predictions,
f"{args.name}_submit_" + str(epoch) + ".csv")
def main(args):
# 1. prepare data & models
train_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
Cutout(10),
RandomBlur(),
TransformByKeys(transforms.ToPILImage(), ("image",)),
TransformByKeys(transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.02), ("image",)),
TransformByKeys(transforms.ToTensor(), ("image",)),
TransformByKeys(transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ("image",)),
])
val_transforms = transforms.Compose([
ScaleMinSideToSize((CROP_SIZE, CROP_SIZE)),
CropCenter(CROP_SIZE),
TransformByKeys(transforms.ToPILImage(), ("image",)),
TransformByKeys(transforms.ToTensor(), ("image",)),
TransformByKeys(transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ("image",)),
])
print("Creating model...")
device = torch.device("cuda: 0") if args.gpu else torch.device("cpu")
model = models.resnet50(pretrained=True)
if args.freeze > 0:
ct = 0
for child in model.children():
ct += 1
if ct <= args.freeze + 4:
for param in child.parameters():
param.requires_grad = False
model.fc = nn.Linear(model.fc.in_features, 2 * NUM_PTS, bias=True)
startEpoch = args.cont
if startEpoch > 0:
with open(f"{args.name}_best_{startEpoch}.pth", "rb") as fp:
best_state_dict = torch.load(fp, map_location="cpu")
model.load_state_dict(best_state_dict)
model.to(device)
if args.test:
val_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'), val_transforms, split="train")
val_dataloader = data.DataLoader(val_dataset, batch_size=args.batch_size, num_workers=4, pin_memory=True,
shuffle=False, drop_last=False)
val_loss_fn = fnn.mse_loss
val_full = validate_full(model, val_dataloader, val_loss_fn, device=device)
res = dict(sorted(val_full.items(), key=lambda x: x[1], reverse=True)[:100])
js = json.dumps(res)
with open(f"{args.name}.json", "w") as f:
f.write(js)
print(res)
return
if not args.predict:
print("Reading data...")
train_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'), train_transforms, split="train")
train_dataloader = data.DataLoader(train_dataset, batch_size=args.batch_size, num_workers=4, pin_memory=True,
shuffle=True, drop_last=True)
val_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'train'), val_transforms, split="val")
val_dataloader = data.DataLoader(val_dataset, batch_size=args.batch_size, num_workers=4, pin_memory=True,
shuffle=False, drop_last=False)
optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=0.0001, nesterov=True)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=2, verbose=True)
train_loss_fn = nn.SmoothL1Loss(reduction="mean")
val_loss_fn = fnn.mse_loss
# 2. train & validate
print("Ready for training...")
best_val_loss = np.inf
for epoch in range(startEpoch, args.epochs):
train_loss = train(model, train_dataloader, train_loss_fn, optimizer, device=device)
val_loss = validate(model, val_dataloader, val_loss_fn, device=device)
scheduler.step(val_loss)
print("Epoch #{:2}:\ttrain loss: {:.5f}\tval loss: {:.5f}".format(epoch, train_loss, val_loss))
with open(f"{args.name}_res.txt", 'a+') as file:
file.write("Epoch #{:2}:\ttrain loss: {:.5f}\tval loss: {:.5f}\n".format(epoch, train_loss, val_loss))
if val_loss < best_val_loss:
best_val_loss = val_loss
with open(f"{args.name}_best.pth", "wb") as fp:
torch.save(model.state_dict(), fp)
if epoch > startEpoch and epoch % 5 == 0:
best_val_loss = val_loss
with open(f"{args.name}_best_{epoch}.pth", "wb") as fp:
torch.save(model.state_dict(), fp)
# 3. predict
test_dataset = ThousandLandmarksDataset(os.path.join(args.data, 'test'), val_transforms, split="test")
test_dataloader = data.DataLoader(test_dataset, batch_size=args.batch_size, num_workers=4, pin_memory=True,
shuffle=False, drop_last=False)
with open(f"{args.name}_best.pth", "rb") as fp:
best_state_dict = torch.load(fp, map_location="cpu")
model.load_state_dict(best_state_dict)
for layer in model.modules():
layer.eval()
test_predictions = predict(model, test_dataloader, device)
with open(f"{args.name}_test_predictions.pkl", "wb") as fp:
pickle.dump({"image_names": test_dataset.image_names,
"landmarks": test_predictions}, fp)
create_submission(args.data, test_predictions, f"{args.name}_submit.csv")
