TransUnet代码复现,使用cpu测试test部分

小菜鸟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()

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值