小菜鸟MOMOMO的transunet复现记录
修改的地方:
1.:把test.py文件中所有的../改成./,../是上一级文件夹,./是当前文件夹的意思
2. 原来的test文件是在gpu环境下运行的,我需要改成在cpu环境下运行,修改的部分如下:
①将设备更改为CPU
在这个截图的地方加上device = torch.device('cpu')
②将模型加载到CPU上
找到和下面相似的那句话,没记错的话它是括号里面的最后那个地方有个.cuda,改成下面这样就行
net = ViT_seg(config_vit, img_size=args.img_size, num_classes=config_vit.n_classes) # 修改后把.cuda删除了
③将输入数据转换为CPU张量:
image, label, case_name = sampled_batch["image"], sampled_batch["label"], sampled_batch['case_name'][0]
修改为:
image, label, case_name = sampled_batch["image"].to(device), sampled_batch["label"].to(device), sampled_batch['case_name'][0]
④
将测试数据加载器的num_workers
参数更改为0:
将以下行代码:
testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=1)
修改为:
testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=0) ####改了,num_worker之前是1
⑤将模型加载的torch.load
函数的map_location
参数设置为device
把
net = ViT_seg(config_vit, img_size=args.img_size,num_classes=config_vit.n_classes).cuda()
改成
net = ViT_seg(config_vit, img_size=args.img_size, num_classes=config_vit.n_classes) # 修改后把.cuda删除了
这么运行了一下还是报错
File "G:\anaconda\envs\pytorch3.7\lib\site-packages\torch\cuda\__init__.py", line 221, in _lazy_init
raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled
是因为utils.py文件当中的gpu设置没有改成cpu
找到def test_sigle_volume
把这两行当中的.cuda改成.cpu就可以了,下面展示的是改后
(1)
input = torch.from_numpy(slice).unsqueeze(0).unsqueeze(0).float().cpu()
(2)
input = torch.from_numpy(image).unsqueeze( 0).unsqueeze(0).float().cpu()