args = parser.parse_args()
#Set random seed for Pytorch and Numpy for reproducibility
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.manual_seed(args.seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(args.seed)
np.random.seed(args.seed)
random.seed(args.seed)
if args.gan == "True":
gan_weight = config.D_WEIGHT
else:
gan_weight = 0.
model = HNNNet(pretrained=True, class_number=2)
if config.D_MULTIPLY:
dnet = DNet(input_dim=3, output_dim=1, input_size=config.PATCH_SIZE)
else:
dnet = DNet(input_dim=4, output_dim=1, input_size=config.PATCH_SIZE)
g_optimizer = optim.SGD(model.parameters(),
lr=config.G_LEARNING_RATE,
momentum=0.9,
weight_decay=0.0005)
d_optimizer = optim.SGD(dnet.parameters(),
lr=config.D_LEARNING_RATE,
momentum=0.9,
weight_decay=0.0005)
resume = config.RESUME_MODEL
'step': tot_step_count,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(state,
os.path.join(dir_checkpoint, 'model_{}.pth.tar'.format(epoch + 1)))
print('Checkpoint {} saved !'.format(epoch + 1))
logger.image_summary('train_images', [vis_image.cpu().numpy() for vis_image in vis_images], step=tot_step_count)
if __name__ == '__main__':
if net_name == 'unet':
model = UNet(n_channels=3, n_classes=6)
else:
model = HNNNet(pretrained=True, class_number=6)
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']+1
start_step = checkpoint['step']
model.load_state_dict(checkpoint['state_dict'])
print('Model loaded from {}'.format(args.resume))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
start_epoch = 0
start_step = 0
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--seed', type=int, default=1234)
parser.add_argument('--model', type=str)
parser.add_argument('--lesion', type=str)
args = parser.parse_args()
#Set random seed for Pytorch and Numpy for reproducibility
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.manual_seed(args.seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(args.seed)
np.random.seed(args.seed)
random.seed(args.seed)
model = HNNNet(pretrained=True, class_number=2)
resume = args.model
if os.path.isfile(resume):
print("=> loading checkpoint '{}'".format(resume))
checkpoint = torch.load(resume)
start_epoch = checkpoint['epoch'] + 1
start_step = checkpoint['step']
try:
model.load_state_dict(checkpoint['state_dict'])
except:
model.load_state_dict(checkpoint['g_state_dict'])
print('Model loaded from {}'.format(resume))
else:
print("=> no checkpoint found at '{}'".format(resume))