[upd] 1. refactor a lot to simplify the pretraining codes; 2. add tutorial for customizing your own CNN model; 3. update some READMEs
keyu-tian
2 years ago
parent
46f9ad2871
commit
6ffe453fa5
15 changed files with 250 additions and 132 deletions
@ -0,0 +1,89 @@ |
||||
# Copyright (c) ByteDance, Inc. and its affiliates. |
||||
# All rights reserved. |
||||
# |
||||
# This source code is licensed under the license found in the |
||||
# LICENSE file in the root directory of this source tree. |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
from typing import List |
||||
from timm.models.registry import register_model |
||||
|
||||
|
||||
class YourConvNet(nn.Module): |
||||
""" |
||||
This is a template for your custom ConvNet. |
||||
It is required to implement the following three functions: `get_downsample_ratio`, `get_feature_map_channels`, `forward`. |
||||
You can refer to the implementations in `pretrain\models\resnet.py` for an example. |
||||
""" |
||||
|
||||
def get_downsample_ratio(self) -> int: |
||||
""" |
||||
This func would ONLY be used in `SparseEncoder's __init__` (see `pretrain/encoder.py`). |
||||
|
||||
:return: the TOTAL downsample ratio of the ConvNet. |
||||
E.g., for a ResNet-50, this should return 32. |
||||
""" |
||||
raise NotImplementedError |
||||
|
||||
def get_feature_map_channels(self) -> List[int]: |
||||
""" |
||||
This func would ONLY be used in `SparseEncoder's __init__` (see `pretrain/encoder.py`). |
||||
|
||||
:return: a list of the number of channels of each feature map. |
||||
E.g., for a ResNet-50, this should return [256, 512, 1024, 2048]. |
||||
""" |
||||
raise NotImplementedError |
||||
|
||||
def forward(self, inp_bchw: torch.Tensor, hierarchical=False): |
||||
""" |
||||
The forward with `hierarchical=True` would ONLY be used in `SparseEncoder.forward` (see `pretrain/encoder.py`). |
||||
|
||||
:param inp_bchw: input image tensor, shape: (batch_size, channels, height, width). |
||||
:param hierarchical: return the logits (not hierarchical), or the feature maps (hierarchical). |
||||
:return: |
||||
- hierarchical == False: return the logits of the classification task, shape: (batch_size, num_classes). |
||||
- hierarchical == True: return a list of all feature maps, which should have the same length as the return value of `get_feature_map_channels`. |
||||
E.g., for a ResNet-50, it should return a list [1st_feat_map, 2nd_feat_map, 3rd_feat_map, 4th_feat_map]. |
||||
for an input size of 224, the shapes are [(B, 256, 56, 56), (B, 512, 28, 28), (B, 1024, 14, 14), (B, 2048, 7, 7)] |
||||
""" |
||||
raise NotImplementedError |
||||
|
||||
|
||||
@register_model |
||||
def your_convnet_small(pretrained=False, **kwargs): |
||||
raise NotImplementedError |
||||
return YourConvNet(**kwargs) |
||||
|
||||
|
||||
@torch.no_grad() |
||||
def convnet_test(): |
||||
from timm.models import create_model |
||||
cnn = create_model('your_convnet_small') |
||||
print('get_downsample_ratio:', cnn.get_downsample_ratio()) |
||||
print('get_feature_map_channels:', cnn.get_feature_map_channels()) |
||||
|
||||
downsample_ratio = cnn.get_downsample_ratio() |
||||
feature_map_channels = cnn.get_feature_map_channels() |
||||
|
||||
# check the forward function |
||||
B, C, H, W = 4, 3, 224, 224 |
||||
inp = torch.rand(B, C, H, W) |
||||
feats = cnn(inp, hierarchical=True) |
||||
assert isinstance(feats, list) |
||||
assert len(feats) == len(feature_map_channels) |
||||
print([tuple(t.shape) for t in feats]) |
||||
|
||||
# check the downsample ratio |
||||
feats = cnn(inp, hierarchical=True) |
||||
assert feats[-1].shape[-2] == H // downsample_ratio |
||||
assert feats[-1].shape[-1] == W // downsample_ratio |
||||
|
||||
# check the channel number |
||||
for feat, ch in zip(feats, feature_map_channels): |
||||
assert feat.ndim == 4 |
||||
assert feat.shape[1] == ch |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
convnet_test() |
||||
Loading…
Reference in new issue