if file_classes[i, 1:].max():
train_idxs.append(i)
train_idxs = np.asarray(train_idxs)
steps_per_epoch = len(train_idxs) // batch_size
validation_steps = len(val_idxs) // val_batch_size
print('steps_per_epoch', steps_per_epoch, 'validation_steps', validation_steps)
data_train = TrainData(train_idxs)
val_train = ValData(val_idxs)
train_data_loader = DataLoader(data_train, batch_size=batch_size, num_workers=6, shuffle=True, pin_memory=False, drop_last=True)
val_data_loader = DataLoader(val_train, batch_size=val_batch_size, num_workers=6, shuffle=False, pin_memory=False)
model = Res34_Unet_Double().cuda()
params = model.parameters()
optimizer = AdamW(params, lr=0.000008, weight_decay=1e-6)
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[1, 2, 3, 4, 5, 7, 9, 11, 17, 23, 29, 33, 47, 50, 60, 70, 90, 110, 130, 150, 170, 180, 190], gamma=0.5)
model = nn.DataParallel(model).cuda()
snap_to_load = 'res34_cls2_{}_0_best'.format(seed)
print("=> loading checkpoint '{}'".format(snap_to_load))
checkpoint = torch.load(path.join(models_folder, snap_to_load), map_location='cpu')
loaded_dict = checkpoint['state_dict']
seed = int(sys.argv[1])
pre_file = sys.argv[2]
post_file = sys.argv[3]
loc_pred_file = sys.argv[4]
cls_pred_file = sys.argv[5]
pred_folder = 'res34cls2_{}_tuned'.format(seed)
makedirs(pred_folder, exist_ok=True)
models = []
snap_to_load = 'res34_cls2_{}_tuned_best'.format(seed)
model = Res34_Unet_Double(pretrained=None)
model = nn.DataParallel(model)
print("=> loading checkpoint '{}'".format(snap_to_load))
checkpoint = torch.load(path.join(models_folder, snap_to_load),
map_location='cpu')
loaded_dict = checkpoint['state_dict']
sd = model.state_dict()
for k in model.state_dict():
if k in loaded_dict and sd[k].size() == loaded_dict[k].size():
sd[k] = loaded_dict[k]
loaded_dict = sd
model.load_state_dict(loaded_dict)
print("loaded checkpoint '{}' (epoch {}, best_score {})".format(
snap_to_load, checkpoint['epoch'], checkpoint['best_score']))