实战PyTorch（二+）：宝可梦自定义数据集之迁移学习

4 steps

1. Load data
2. Build model
3. Train and Test
4. Transfer Learning

• Load data

Image Resize：224x224 for ResNet18

• Build model

1. Inherit from base class：在ResNet18结构上做修改。
2. Define forward graph

class ResBlk(nn.Module):
    def __init__(self, ch_in, ch_out, stride=1):
        super(ResBlk, self).__init__()
        self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(ch_out)
        self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(ch_out)
        '''extra [b, ch_in, h, w] => [b, ch_out, h, w]'''
        self.extra = nn.Sequential()
        if ch_out != ch_in:
            self.extra = nn.Sequential(
                nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=stride),
                nn.BatchNorm2d(ch_out)
            )
    def forward(self, x):
        """param x: [b, ch, h, w]"""
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        # short cut.
        # extra module: [b, ch_in, h, w] => [b, ch_out, h, w]
        # element-wise add:
        out = self.extra(x) + out
        out = F.relu(out)
        return out

class ResNet18(nn.Module):
    def __init__(self, num_class):
        super(ResNet18, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=3, padding=0),
            nn.BatchNorm2d(16)
        )
        # followed 4 blocks
        # [b, 16, h, w] => [b, 32, h ,w]
        self.blk1 = ResBlk(16, 32, stride=3)
        # [b, 32, h, w] => [b, 64, h, w]
        self.blk2 = ResBlk(32, 64, stride=3)
        # # [b, 64, h, w] => [b, 128, h, w]
        self.blk3 = ResBlk(64, 128, stride=2)
        # # [b, 128, h, w] => [b, 256, h, w]
        self.blk4 = ResBlk(128, 256, stride=2)
        # [b, 256, 3, 3] 全连接层输出
        self.outlayer = nn.Linear(256*3*3, num_class)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        # [b, 64, h, w] => [b, 1024, h, w]
        x = self.blk1(x)
        x = self.blk2(x)
        x = self.blk3(x)
        x = self.blk4(x)
        # print(x.shape)
        x = x.view(x.size(0), -1)
        x = self.outlayer(x)
        return x

• Train and Test

总体思想：

选最好的状态保存起来做测试。

 

1.首先读取数据集

train_db = Pokemon('pokemon', 224, mode='train')
val_db = Pokemon('pokemon', 224, mode='val')
test_db = Pokemon('pokemon', 224, mode='test')
train_loader = DataLoader(train_db, batch_size=batchsz, shuffle=True,num_workers=4)
val_loader = DataLoader(val_db, batch_size=batchsz, num_workers=2)
test_loader = DataLoader(test_db, batch_size=batchsz, num_workers=2)

2.让验证集选最好的状态保存起来，让模型训练效果更好一点，以便以后做测试。

def evalute(model, loader):
    model.eval()
    correct = 0
    total = len(loader.dataset)
    for x,y in loader:
        x,y = x.to(device), y.to(device)
        with torch.no_grad():
            logits = model(x)
            pred = logits.argmax(dim=1)
        correct += torch.eq(pred, y).sum().float().item()
    return correct / total

3.main()

def main():
    model = ResNet18(5).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criteon = nn.CrossEntropyLoss()

    best_acc, best_epoch = 0, 0
    global_step = 0
    viz.line([0], [-1], win='loss', opts=dict(title='loss'))
    viz.line([0], [-1], win='val_acc', opts=dict(title='val_acc'))
'''模型的训练及优化'''
    for epoch in range(epochs):
        for step, (x,y) in enumerate(train_loader):
            x, y = x.to(device), y.to(device)  # x: [b, 3, 224, 224], y: [b]
            '''模型训练及优化'''
            model.train()
            logits = model(x)
            loss = criteon(logits, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            viz.line([loss.item()], [global_step], win='loss', update='append')
            global_step += 1
            '''对比并保存训练每回合最优的模型参数'''
        if epoch % 1 == 0:
            val_acc = evalute(model, val_loader)
            if val_acc> best_acc:
                best_epoch = epoch
                best_acc = val_acc
                torch.save(model.state_dict(), 'best.mdl')
            viz.line([val_acc], [global_step], win='val_acc', update='append')

    '''输出acc结果'''
    print('best acc:', best_acc, 'best epoch:', best_epoch)
'''测试'''
    model.load_state_dict(torch.load('best.mdl'))
    print('loaded from ckpt!')
    test_acc = evalute(model, test_loader)
    print('test acc:', test_acc)

上面模型的测试结果在0.85左右，不是特别理想，想提升的话可以使用迁移学习。

 

• Transfer Learning

• Others

'''测试函数'''
    '''ResBlk测试'''
    blk = ResBlk(64, 128) #ch_out[64]=>ch_out[128]
    tmp = torch.randn(2, 64, 224, 224) #输入两张，通道为64，宽高为224的图片。
    out = blk(tmp)
    print('block:', out.shape)#[2, 128, 224, 224]
    '''Resnet18测试'''
    model = ResNet18(5)
    tmp = torch.randn(2, 3, 224, 224)
    out = model(tmp)
    print('resnet:', out.shape) #[2,5]
    '''总参数量'''
    p = sum(map(lambda p:p.numel(), model.parameters())) 
    print('parameters size:', p) #1234885