def main():
# settting
sets = parse_opts()
model, parameters = generate_model(sets)
#debug
# print(model)
# print(f'type(model):{type(model)}')
# print(f'type(parameters):{type(parameters)}')
# print(f'type(model.parameters):{type(model.parameters)}')
# freeze the pre-trained convolutional network parameters
# for param in parameters['base_parameters']:
# param.requires_grad = False
# ----------------------------------------------------
# Debug:
# for key in parameters.keys():
# _list = parameters[key]
# for i, item in enumerate(_list):
# print(item.shape)
# print(item.requires_grad)
# criterion
criterion = CrossEntropyLoss()
# optimizer
params = [{
'params': parameters['base_parameters'],
'lr': sets.learning_rate
}, {
'params': parameters['new_parameters'],
'lr': sets.learning_rate * 100
}]
optimizer = optim.Adam(
params,
lr=0.1,
)
# train loader
train_dataset = ABIDE1(sets, train=True)
train_loader = DataLoader(train_dataset,
batch_size=8,
shuffle=True,
num_workers=1,
drop_last=True)
# test loader
# test_dataset = ABIDE1(sets, train=False)
# test_loader = DataLoader(test_dataset, num_workers=1)
train(train_loader, model, optimizer, criterion)
torch.save(
{
'ecpoch': epoch,
'batch_id': batch_id,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, model_save_path)
print('Finished training')
if sets.ci_test:
exit()
if __name__ == '__main__':
# settting
sets = parse_opts()
if sets.ci_test:
sets.img_list = './toy_data/test_ci.txt'
sets.n_epochs = 1
sets.no_cuda = True
sets.data_root = './toy_data'
sets.pretrain_path = ''
sets.num_workers = 0
sets.model_depth = 10
sets.resnet_shortcut = 'A'
sets.input_D = 14
sets.input_H = 28
sets.input_W = 28
# getting model
torch.manual_seed(sets.manual_seed)
dices[idx, :] = seg_eval(pred_total[idx], y_total[idx],
range(self.hparams.n_seg_classes))
print("avg_loss(val) = ", avg_loss)
print("dice_each_class(val) = ", dices.mean(axis=0))
self.logger.experiment.add_text("dice_each_class",
str(dices.mean(axis=0)),
self.current_epoch)
# print(self.trainer.lr_schedulers.get_lr())
mDice = dices.mean()
logs = {"val_loss": avg_loss, "mDice": torch.tensor(mDice)}
return {'log': logs}
if __name__ == '__main__':
# settings
hparams = parse_opts()
checkpoint_callback = ModelCheckpoint(filepath=None,
monitor='mDice',
save_top_k=1,
verbose=True,
mode='max')
lr_logger = LearningRateLogger()
if hparams.phase == "test":
pretrained_model = BiMaskSeg.load_from_checkpoint(
checkpoint_path=
"trails/logs/train_lightning/lightning_logs/version_0/checkpoints/epoch=29.ckpt",
hparams_file=
"trails/logs/train_lightning/lightning_logs/version_0/hparams.yaml"
log.info('Save checkpoints: epoch = {}, batch_id = {}'.format(epoch, batch_id))
torch.save({
'epoch': epoch,
'batch_id': batch_id,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()},
model_save_path)
scheduler.step()
print('Finished training')
if __name__ == '__main__':
opt = parse_opts()
str_ids = opt.gpu_ids.split(',')
opt.gpu_ids = []
for str_id in str_ids:
idx = int(str_id)
if idx >= 0:
opt.gpu_ids.append(idx)
torch.manual_seed(opt.seed)
dataloader = DataLoader(opt)
model = RegresserModel(opt)
if not opt.no_cuda:
if len(opt.gpu_ids) > 1:
