def main(args, num_fold=0):
# 模型选择
model = seg_model(args)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# if args.mode == "train" and num_fold <= 1:
# torchsummary.summary(model, (3, args.crop_size[0], args.crop_size[1])) # #输出网络结构和参数量
print(f' [network: {args.network} device: {device}]')
if args.mode == "train":
try:
train(model, device, args, num_fold=num_fold)
except KeyboardInterrupt:
torch.save(
model.state_dict(),
os.path.join(args.checkpoint_dir[num_fold], 'INTERRUPTED.pth'))
print('Saved interrupt')
try:
sys.exit(0)
except SystemExit:
os._exit(0)
elif args.mode == "test":
if args.k_fold is not None:
return test(model, device, args, num_fold=num_fold)
else:
test(model, device, args, num_fold=num_fold)
else:
raise NotImplementedError
def main(args, num_fold=0):
# 模型选择
model = seg_model(args)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# if args.mode == "train" and num_fold <= 1:
# torchsummary.summary(model, (3, args.crop_size[0], args.crop_size[1])) # #输出网络结构和参数量
print(f' [network: {args.network} device: {device}]')
if args.mode == "train":
train(model, device, args, num_fold=num_fold)
elif args.mode == "test":
if args.k_fold is not None:
return test(model, device, args, num_fold=num_fold)
else:
test(model, device, args, num_fold=num_fold)
else:
raise NotImplementedError
pred_mask = pred_mask.squeeze().numpy()
pred_mask = Image.fromarray(np.uint16(pred_mask + 1) * 100)
pred_mask.save(
os.path.join(pred_dir[num_fold],
file_name[i] + ".png"))
pbar.update(image.shape[0])
if __name__ == "__main__":
torch.cuda.empty_cache()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
args = basic_setting()
# 模型选择
model = seg_model(args)
pred_dir = [os.path.join(args.dir, pred_dir)]
map_dir = [os.path.join(args.dir, map_dir)]
if not os.path.exists(pred_dir[0]):
os.mkdir(pred_dir[0])
if not os.path.exists(map_dir[0]):
os.mkdir(map_dir[0])
if args.k_fold == None:
model.to(device)
model_dir = os.path.join(args.checkpoint_dir[0],
f'CP_epoch{model_CPepoch}.pth') # 最后一个epoch模型
model.load_state_dict(torch.load(model_dir, map_location=device))
print("model loaded!")