TransUnet使用自己的数据集复现求改进（抛砖引玉交流贴）

        在使用过多个基于transformer的开源网络过后，包括segformer、objectformer、swinunet、transfuse等，发现加载官方公布的预训练权重的transunet在自己的任务上表现最好（我的方向是图像复制篡改检测，就是一种二分类语义分割），在不做任何修改的情况仅重新成功训练就有远超普通全卷积网络的能力，（这里仅指transunet的R50-ViT-B_16这种卷积后接transformer的形式，而同样加载预训练权重的纯vit模型表现就不尽如人意了，很不理解？）

        先帮还没入手的的兄弟解决些复现的小问题。

        官方code：GitHub - Beckschen/TransUNet: This repository includes the official project of TransUNet, presented in our paper: TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation.

里面还有所有模型预训练权重和一些训练数据集的下载链接，就在readme哪儿。

        在自己的数据集上复现还是比较简单的，基本上重写一下dataloader就行了，我直接就贴我的代码了，改的另一位博主如雾如电的，稍微借鉴了原始的数据读取但我没有加任何图像预处理。

def own_data_loader(img_path, mask_path):
    img = cv2.imread(img_path)
    img = cv2.resize(img, (224,224), interpolation = cv2.INTER_NEAREST)
    mask = cv2.imread(mask_path, 0)

    mask = cv2.resize(mask, (224,224), interpolation = cv2.INTER_NEAREST)

    # img = randomHueSaturationValue(img,
    #                                hue_shift_limit=(-30, 30),
    #                                sat_shift_limit=(-5, 5),
    #                                val_shift_limit=(-15, 15))
    #
    # img, mask = randomShiftScaleRotate(img, mask,
    #                                    shift_limit=(-0.1, 0.1),
    #                                    scale_limit=(-0.1, 0.1),
    #                                    aspect_limit=(-0.1, 0.1),
    #                                    rotate_limit=(-0, 0))
    # img, mask = randomHorizontalFlip(img, mask)
    # img, mask = randomVerticleFlip(img, mask)
    # img, mask = randomRotate90(img, mask)

    # mask = np.expand_dims(mask, axis=2)

    img = np.array(img, np.float32) / 255.0 * 3.2 - 1.6
    # img = np.array(img, np.float32) / 255.0
    # mask = np.array(mask, np.float32)
    mask = np.array(mask, np.float32) / 255.0
    mask[mask >= 0.5] = 1
    mask[mask < 0.5] = 0
    # mask = np.squeeze(mask,axis=2)
    # pyplot.imshow(mask)
    # pyplot.show()

    img = np.array(img, np.float32).transpose(2, 0, 1)
    # mask = np.array(mask, np.float32).transpose(2, 0, 1)
    return img, mask

def read_own_data(root_path, mode='train'):
    images = []
    masks = []

    image_root = os.path.join(root_path+'/images'+ '/train')

    gt_root = os.path.join(root_path+ "/annotations"+ '/train_png')

    for image_name in os.listdir(gt_root):
        label_path = os.path.join(gt_root, image_name)
        masks.append(label_path)
    for image_name in os.listdir(image_root):
        image_path = os.path.join(image_root, image_name)
        images.append(image_path)


    return images, masks

class ImageFolder(Dataset):
    def __init__(self,root_path, mode='train'):
        self.root = root_path
        self.mode = mode
        self.images, self.labels = read_own_data(self.root, self.mode)

    def __getitem__(self, index):
        if self.mode == 'test':
            img, mask = own_data_test_loader(self.images[index], self.labels[index])
        else:
            img, mask = own_data_loader(self.images[index], self.labels[index])
            # img = torch.Tensor(img)
            # mask = torch.Tensor(mask)
        return img, mask

    def __len__(self):
        # assert len(self.images) == len(self.labels), 'The number of images must be equal to labels'
        return len(self.images)

）

调用就一行，在trainer.py好像还要改一点点

db_train = ImageFolder( args.root_path,mode='train')

iterator = tqdm(range(max_epoch), ncols=70)
for epoch_num in iterator:

    progress_bar = tqdm(trainloader)
    for image_batch, label_batch in trainloader:
        # image_batch, label_batch = sampled_batch['image'], sampled_batch['label']
        image_batch, label_batch = image_batch.cuda(), label_batch.cuda()
        # print(image_batch.shape, "333")
        outputs = model(image_batch)
        # outputs = nn.Sigmoid()(outputs)
        # print(outputs) #torch.Size([6, 2, 224, 224])
        # print(label_batch.shape) #torch.Size([6, 1, 224, 224])

        # 这里的ce_loss = CrossEntropyLoss()常用于多分类，换成BCELoss
        # loss_ce = ce_loss(outputs, label_batch[:].long())
        # loss_dice = dice_loss(outputs, label_batch, softmax=True)
        # loss = 0.4 * loss_ce + 0.6 * loss_dice
        outputs = torch.squeeze(outputs)
        label_batch = torch.squeeze(label_batch)



        # loss_ce = bce_loss(outputs, label_batch)
        # loss_dice = dice_loss(outputs, label_batch)
        # loss = 0.4 * loss_ce + 0.6 * loss_dice

        loss = bce_loss(outputs, label_batch)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr_

        iter_num = iter_num + 1

        progress_bar.set_description(
            ' Epoch: {}/{}. Iteration: {}/{}. Mini loss: {:.5f}  '.format(
                epoch_num, 30, iter_num , max_iterations, loss.item(),))

我这里也改了损失函数，用惯了Bce了，不过大家还是先用一下原始损失函数再根据自己任务调整。

其它基本就是改一些参数就行了，到了要加载预训练权重的时候可以发现，公布的权重并不是整个网络的，只包括resnet50三层和12个transformer层，不包括后续四个decoder层。这里就是我迫切想和诸位交流并咨询的一点，我做了大量实验表明，只有全部加载提供的预训练权重才能得到一个差不多的精度。不论你是只加载resnet50三层或者12个transformer层或者将resnet50三层换成同样有imagenet预训练权重的vgg等其它网络都会极大影响性能。

        而且在resnet50与12个transformer之间新增一些模块或操作也会同样影响性能，不过对于在decoder部分做出的改动对于性能影响不大。

        现在的问题就是resnet50与12个transformer深度绑定而且必须加载预训练权重才能保证性能，这就导致你不能对这两部分做出任何改动。所以想请教大家如何才能获得这篇论文的一个原始预训练数据，直接把imagenet拿过来做图像分类就行吗？

        因为想靠transformer再混一篇出来，所以还想请教大家还有哪些带有imagenet预训练权重的transformer结构？因为以我的经验来看，哪些没带预训练权重的former完全没办法成功训练，同样的参数下它的损失很难收敛。