def main():
# option flags
FLG = train_args()
# torch setting
device = torch.device('cuda:{}'.format(FLG.devices[0]))
torch.backends.cudnn.benchmark = True
torch.cuda.set_device(FLG.devices[0])
# create summary and report the option
visenv = FLG.model
summary = Summary(port=39199, env=visenv)
summary.viz.text(argument_report(FLG, end='<br>'),
win='report' + str(FLG.running_fold))
train_report = ScoreReport()
valid_report = ScoreReport()
timer = SimpleTimer()
fold_str = 'fold' + str(FLG.running_fold)
best_score = dict(epoch=0, loss=1e+100, accuracy=0)
#### create dataset ###
# kfold split
target_dict = np.load(pjoin(FLG.data_root, 'target_dict.pkl'))
trainblock, validblock, ratio = fold_split(
FLG.fold, FLG.running_fold, FLG.labels,
np.load(pjoin(FLG.data_root, 'subject_indices.npy')), target_dict)
def _dataset(block, transform):
return ADNIDataset(FLG.labels,
pjoin(FLG.data_root, FLG.modal),
block,
target_dict,
transform=transform)
# create train set
trainset = _dataset(trainblock, transform_presets(FLG.augmentation))
# create normal valid set
validset = _dataset(
validblock,
transform_presets('nine crop' if FLG.augmentation ==
'random crop' else 'no augmentation'))
# each loader
trainloader = DataLoader(trainset,
batch_size=FLG.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True)
validloader = DataLoader(validset, num_workers=4, pin_memory=True)
# data check
# for image, _ in trainloader:
# summary.image3d('asdf', image)
# create model
def kaiming_init(tensor):
return kaiming_normal_(tensor, mode='fan_out', nonlinearity='relu')
if 'plane' in FLG.model:
model = Plane(len(FLG.labels),
name=FLG.model,
weights_initializer=kaiming_init)
elif 'resnet11' in FLG.model:
model = resnet11(len(FLG.labels),
FLG.model,
weights_initializer=kaiming_init)
elif 'resnet19' in FLG.model:
model = resnet19(len(FLG.labels),
FLG.model,
weights_initializer=kaiming_init)
elif 'resnet35' in FLG.model:
model = resnet35(len(FLG.labels),
FLG.model,
weights_initializer=kaiming_init)
elif 'resnet51' in FLG.model:
model = resnet51(len(FLG.labels),
FLG.model,
weights_initializer=kaiming_init)
else:
raise NotImplementedError(FLG.model)
print_model_parameters(model)
model = torch.nn.DataParallel(model, FLG.devices)
model.to(device)
# criterion
train_criterion = torch.nn.CrossEntropyLoss(weight=torch.Tensor(
list(map(lambda x: x * 2, reversed(ratio))))).to(device)
valid_criterion = torch.nn.CrossEntropyLoss().to(device)
# TODO resume
# optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=FLG.lr,
weight_decay=FLG.l2_decay)
# scheduler
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, FLG.lr_gamma)
start_epoch = 0
global_step = start_epoch * len(trainloader)
pbar = None
for epoch in range(1, FLG.max_epoch + 1):
timer.tic()
scheduler.step()
summary.scalar('lr',
fold_str,
epoch - 1,
optimizer.param_groups[0]['lr'],
ytickmin=0,
ytickmax=FLG.lr)
# train()
torch.set_grad_enabled(True)
model.train(True)
train_report.clear()
if pbar is None:
pbar = tqdm(total=len(trainloader) * FLG.validation_term,
desc='Epoch {:<3}-{:>3} train'.format(
epoch, epoch + FLG.validation_term - 1))
for images, targets in trainloader:
images = images.cuda(device, non_blocking=True)
targets = targets.cuda(device, non_blocking=True)
optimizer.zero_grad()
outputs = model(images)
loss = train_criterion(outputs, targets)
loss.backward()
optimizer.step()
train_report.update_true(targets)
train_report.update_score(F.softmax(outputs, dim=1))
summary.scalar('loss',
'train ' + fold_str,
global_step / len(trainloader),
loss.item(),
ytickmin=0,
ytickmax=1)
pbar.update()
global_step += 1
if epoch % FLG.validation_term != 0:
timer.toc()
continue
pbar.close()
# valid()
torch.set_grad_enabled(False)
model.eval()
valid_report.clear()
pbar = tqdm(total=len(validloader),
desc='Epoch {:>3} valid'.format(epoch))
for images, targets in validloader:
true = targets
npatchs = 1
if len(images.shape) == 6:
_, npatchs, c, x, y, z = images.shape
images = images.view(-1, c, x, y, z)
targets = torch.cat([targets
for _ in range(npatchs)]).squeeze()
images = images.cuda(device, non_blocking=True)
targets = targets.cuda(device, non_blocking=True)
output = model(images)
loss = valid_criterion(output, targets)
valid_report.loss += loss.item()
if npatchs == 1:
score = F.softmax(output, dim=1)
else:
score = torch.mean(F.softmax(output, dim=1),
dim=0,
keepdim=True)
valid_report.update_true(true)
valid_report.update_score(score)
pbar.update()
pbar.close()
# report
vloss = valid_report.loss / len(validloader)
summary.scalar('accuracy',
'train ' + fold_str,
epoch,
train_report.accuracy,
ytickmin=-0.05,
ytickmax=1.05)
summary.scalar('loss',
'valid ' + fold_str,
epoch,
vloss,
ytickmin=0,
ytickmax=0.8)
summary.scalar('accuracy',
'valid ' + fold_str,
epoch,
valid_report.accuracy,
ytickmin=-0.05,
ytickmax=1.05)
is_best = False
if best_score['loss'] > vloss:
best_score['loss'] = vloss
best_score['epoch'] = epoch
best_score['accuracy'] = valid_report.accuracy
is_best = True
print('Best Epoch {}: validation loss {} accuracy {}'.format(
best_score['epoch'], best_score['loss'], best_score['accuracy']))
# save
if isinstance(model, torch.nn.DataParallel):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint(
dict(epoch=epoch,
best_score=best_score,
state_dict=state_dict,
optimizer_state_dict=optimizer.state_dict()),
FLG.checkpoint_root, FLG.running_fold, FLG.model, is_best)
pbar = None
timer.toc()
print('Time elapse {}h {}m {}s'.format(*timer.total()))
def test(FLG):
device = torch.device('cuda:{}'.format(FLG.devices[0]))
torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True
torch.cuda.set_device(FLG.devices[0])
report = [ScoreReport() for _ in range(FLG.fold)]
overall_report = ScoreReport()
target_dict = np.load(pjoin(FLG.data_root, 'target_dict.pkl'))
with open(FLG.model + '_stat.pkl', 'rb') as f:
stat = pickle.load(f)
summary = Summary(port=10001, env=str(FLG.model) + 'CAM')
class Feature(object):
def __init__(self):
self.blob = None
def capture(self, blob):
self.blob = blob
if 'plane' in FLG.model:
model = Plane(len(FLG.labels), name=FLG.model)
elif 'resnet11' in FLG.model:
model = resnet11(len(FLG.labels), FLG.model)
elif 'resnet19' in FLG.model:
model = resnet19(len(FLG.labels), FLG.model)
elif 'resnet35' in FLG.model:
model = resnet35(len(FLG.labels), FLG.model)
elif 'resnet51' in FLG.model:
model = resnet51(len(FLG.labels), FLG.model)
else:
raise NotImplementedError(FLG.model)
model.to(device)
ad_h = []
nl_h = []
adcams = np.zeros((4, 3, 112, 144, 112), dtype="f8")
nlcams = np.zeros((4, 3, 112, 144, 112), dtype="f8")
sb = [9.996e-01, 6.3e-01, 1.001e-01]
for running_fold in range(FLG.fold):
_, validblock, _ = fold_split(
FLG.fold, running_fold, FLG.labels,
np.load(pjoin(FLG.data_root, 'subject_indices.npy')), target_dict)
validset = ADNIDataset(FLG.labels,
pjoin(FLG.data_root, FLG.modal),
validblock,
target_dict,
transform=transform_presets(FLG.augmentation))
validloader = DataLoader(validset, pin_memory=True)
epoch, _ = load_checkpoint(model, FLG.checkpoint_root, running_fold,
FLG.model, None, True)
model.eval()
feature = Feature()
def hook(mod, inp, oup):
return feature.capture(oup.data.cpu().numpy())
_ = model.layer4.register_forward_hook(hook)
fc_weights = model.fc.weight.data.cpu().numpy()
transformer = Compose([CenterCrop((112, 144, 112)), ToFloatTensor()])
im, _ = original_load(validblock, target_dict, transformer, device)
for image, target in validloader:
true = target
npatches = 1
if len(image.shape) == 6:
_, npatches, c, x, y, z = image.shape
image = image.view(-1, c, x, y, z)
target = torch.stack([target
for _ in range(npatches)]).squeeze()
image = image.cuda(device, non_blocking=True)
target = target.cuda(device, non_blocking=True)
output = model(image)
if npatches == 1:
score = F.softmax(output, dim=1)
else:
score = torch.mean(F.softmax(output, dim=1),
dim=0,
keepdim=True)
report[running_fold].update_true(true)
report[running_fold].update_score(score)
overall_report.update_true(true)
overall_report.update_score(score)
print(target)
if FLG.cam:
s = 0
cams = []
if target[0] == 0:
s = score[0][0]
#s = s.cpu().numpy()[()]
cams = adcams
else:
sn = score[0][1]
#s = s.cpu().numpy()[()]
cams = nlcams
if s > sb[0]:
cams[0] = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
cams[0],
s,
num_images=5)
elif s > sb[1]:
cams[1] = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
cams[1],
s,
num_images=5)
elif s > sb[2]:
cams[2] = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
cams[2],
s,
num_images=5)
else:
cams[3] = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
cams[3],
s,
num_images=5)
#ad_h += [s]
#nl_h += [sn]
print('At {}'.format(epoch))
print(
metrics.classification_report(report[running_fold].y_true,
report[running_fold].y_pred,
target_names=FLG.labels,
digits=4))
print('accuracy {}'.format(report[running_fold].accuracy))
#print(np.histogram(ad_h))
#print(np.histogram(nl_h))
print('over all')
print(
metrics.classification_report(overall_report.y_true,
overall_report.y_pred,
target_names=FLG.labels,
digits=4))
print('accuracy {}'.format(overall_report.accuracy))
with open(FLG.model + '_stat.pkl', 'wb') as f:
pickle.dump(report, f, pickle.HIGHEST_PROTOCOL)
def test(FLG):
device = torch.device('cuda:{}'.format(FLG.devices[0]))
torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True
torch.cuda.set_device(FLG.devices[0])
report = [ScoreReport() for _ in range(FLG.fold)]
overall_report = ScoreReport()
target_dict = np.load(pjoin(FLG.data_root, 'target_dict.pkl'))
if 'plane' in FLG.model:
model = Plane(len(FLG.labels), name=FLG.model)
elif 'resnet11' in FLG.model:
model = resnet11(len(FLG.labels), FLG.model)
elif 'resnet19' in FLG.model:
model = resnet19(len(FLG.labels), FLG.model)
elif 'resnet35' in FLG.model:
model = resnet35(len(FLG.labels), FLG.model)
elif 'resnet51' in FLG.model:
model = resnet51(len(FLG.labels), FLG.model)
else:
raise NotImplementedError(FLG.model)
model.to(device)
for running_fold in range(FLG.fold):
_, validblock, _ = fold_split(
FLG.fold, running_fold, FLG.labels,
np.load(pjoin(FLG.data_root, 'subject_indices.npy')), target_dict)
validset = ADNIDataset(FLG.labels,
pjoin(FLG.data_root, FLG.modal),
validblock,
target_dict,
transform=transform_presets(FLG.augmentation))
validloader = DataLoader(validset, pin_memory=True)
epoch, _ = load_checkpoint(model, FLG.checkpoint_root, running_fold,
FLG.model, None, True)
model.eval()
for image, target in validloader:
true = target
npatches = 1
if len(image.shape) == 6:
_, npatches, c, x, y, z = image.shape
image = image.view(-1, c, x, y, z)
target = torch.stack([target
for _ in range(npatches)]).squeeze()
image = image.cuda(device, non_blocking=True)
target = target.cuda(device, non_blocking=True)
output = model(image)
if npatches == 1:
score = F.softmax(output, dim=1)
else:
score = torch.mean(F.softmax(output, dim=1),
dim=0,
keepdim=True)
report[running_fold].update_true(true)
report[running_fold].update_score(score)
overall_report.update_true(true)
overall_report.update_score(score)
print('At {}'.format(epoch))
print(
metrics.classification_report(report[running_fold].y_true,
report[running_fold].y_pred,
target_names=FLG.labels,
digits=4))
print('accuracy {}'.format(report[running_fold].accuracy))
print('over all')
print(
metrics.classification_report(overall_report.y_true,
overall_report.y_pred,
target_names=FLG.labels,
digits=4))
print('accuracy {}'.format(overall_report.accuracy))
with open(FLG.model + '_stat.pkl', 'wb') as f:
pickle.dump(report, f, pickle.HIGHEST_PROTOCOL)
def test_cam():
with open(FLG.model + '_stat.pkl', 'rb') as f:
stat = pickle.load(f)
summary = Summary(port=10001, env=str(FLG.model) + 'CAM')
class Feature(object):
def __init__(self):
self.blob = None
def capture(self, blob):
self.blob = blob
# TODO: create model
device = torch.device('cuda:{}'.format(FLG.devices[0]))
torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True
torch.cuda.set_device(FLG.devices[0])
report = [ScoreReport() for _ in range(FLG.fold)]
target_dict = np.load(pjoin(FLG.data_root, 'target_dict.pkl'))
model = Plane(len(FLG.labels), name=FLG.model)
model.to(device)
transformer = Compose([CenterCrop((112, 144, 112)), ToFloatTensor()])
def original_load(validblock):
originalset = ADNIDataset(FLG.labels,
pjoin(FLG.data_root, 'spm_normalized'),
validblock,
target_dict,
transform=transformer)
originloader = DataLoader(originalset, pin_memory=True)
for image, target in originloader:
if len(image.shape) == 6:
_, npatches, c, x, y, z = image.shape
image = image.view(-1, c, x, y, z)
target = torch.stack([target
for _ in range(npatches)]).squeeze()
image = image.cuda(device, non_blocking=True)
target = target.cuda(device, non_blocking=True)
break
return image, target
hadcams = np.zeros((3, 112, 144, 112), dtype="f8")
madcams = np.zeros((3, 112, 144, 112), dtype="f8")
sadcams = np.zeros((3, 112, 144, 112), dtype="f8")
zadcams = np.zeros((3, 112, 144, 112), dtype="f8")
nlcams = np.zeros((4, 3, 112, 144, 112), dtype="f8")
sb = [4.34444371e-16, 1.67179015e-18, 4.08813312e-23]
#im, _ = original_load(validblock)
for running_fold in range(FLG.fold):
# validset
_, validblock, _ = fold_split(
FLG.fold, running_fold, FLG.labels,
np.load(pjoin(FLG.data_root, 'subject_indices.npy')), target_dict)
validset = ADNIDataset(FLG.labels,
pjoin(FLG.data_root, FLG.modal),
validblock,
target_dict,
transform=transformer)
validloader = DataLoader(validset, pin_memory=True)
load_checkpoint(model,
FLG.checkpoint_root,
running_fold,
FLG.model,
epoch=None,
is_best=True)
model.eval()
feature = Feature()
def hook(mod, inp, oup):
return feature.capture(oup.data.cpu().numpy())
_ = model.layer4.register_forward_hook(hook)
fc_weights = model.fc.weight.data.cpu().numpy()
im, _ = original_load(validblock)
ad_s = []
for image, target in validloader:
true = target
npatches = 1
if len(image.shape) == 6:
_, npatches, c, x, y, z = image.shape
image = image.view(-1, c, x, y, z)
target = torch.stack([target
for _ in range(npatches)]).squeeze()
image = image.cuda(device, non_blocking=True)
target = target.cuda(device, non_blocking=True)
#_ = model(image.view(*image.shape))
output = model(image)
if npatches == 1:
score = F.softmax(output, dim=1)
else:
score = torch.mean(F.softmax(output, dim=1),
dim=0,
keepdim=True)
sa = score[0][1]
#name = 'k'+str(running_fold)
sa = sa.cpu().numpy()[()]
print(score, score.shape)
if true == torch.tensor([0]):
if sa > sb[0]:
hadcams = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
hadcams,
sa,
num_images=5)
elif sa > sb[1]:
madcams = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
madcams,
sa,
num_images=5)
elif sa > sb[2]:
sadcams = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
sadcams,
sa,
num_images=5)
else:
zadcams = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
zadcams,
sa,
num_images=5)
else:
if s > sb[0]:
nlcams[0] = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
nlcams[0],
sr,
num_images=5)
elif sr > sb[1]:
nlcams[1] = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
nlcams[1],
sr,
num_images=5)
elif sr > sb[2]:
nlcams[2] = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
nlcams[2],
sr,
num_images=5)
else:
nlcams[3] = summary.cam3d(FLG.labels[target],
im,
feature.blob,
fc_weights,
target,
nlcams[3],
sr,
num_images=5)
ad_s += [sr]
print('histogram', np.histogram(ad_s))