Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from torchvision.models import resnet | |
| import torch.utils.model_zoo as model_zoo | |
| from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg | |
| model_urls = { | |
| 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', | |
| 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', | |
| 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', | |
| 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', | |
| 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', | |
| } | |
| class ResNet(nn.Module): | |
| def __init__(self, name="resnet50", pretrain=True): | |
| super().__init__() | |
| if name == "resnet50": | |
| base_net = resnet.resnet50(pretrained=False) | |
| elif name == "resnet101": | |
| base_net = resnet.resnet101(pretrained=False) | |
| elif name == "resnet18": | |
| base_net = resnet.resnet18(pretrained=False) | |
| elif name == "resnet34": | |
| base_net = resnet.resnet34(pretrained=False) | |
| else: | |
| print(" base model is not support !") | |
| if pretrain: | |
| print("load the {} weight from ./cache".format(name)) | |
| base_net.load_state_dict(model_zoo.load_url(model_urls[name], model_dir="./cache", | |
| map_location=torch.device(cfg.device)), strict=False) | |
| # print(base_net) | |
| self.stage1 = nn.Sequential( | |
| base_net.conv1, | |
| base_net.bn1, | |
| base_net.relu, | |
| base_net.maxpool | |
| ) | |
| self.stage2 = base_net.layer1 | |
| self.stage3 = base_net.layer2 | |
| self.stage4 = base_net.layer3 | |
| self.stage5 = base_net.layer4 | |
| self.up2 = nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1) | |
| def forward(self, x): | |
| C1 = self.stage1(x) | |
| C2 = self.stage2(C1) | |
| C3 = self.stage3(C2) | |
| C4 = self.stage4(C3) | |
| C5 = self.stage5(C4) | |
| if cfg.scale == 2 or cfg.scale == 1: | |
| # up2 --> 1/2 | |
| C1 = self.up2(C1) | |
| return C1, C2, C3, C4, C5 | |
| if __name__ == '__main__': | |
| import torch | |
| input = torch.randn((4, 3, 512, 512)) | |
| net = ResNet() | |
| C1, C2, C3, C4, C5 = net(input) | |
| print(C1.size()) | |
| print(C2.size()) | |
| print(C3.size()) | |
| print(C4.size()) | |
| print(C5.size()) | |