|
@@ -1,11 +1,11 @@
|
|
|
import torch
|
|
|
-import torchvision.models as tmodels
|
|
|
+from torchvision.models import resnet18, ResNet18_Weights
|
|
|
import torch.nn as nn
|
|
|
|
|
|
class ModelCT(nn.Module):
|
|
|
def __init__(self):
|
|
|
super(ModelCT, self).__init__()
|
|
|
- self.backbone = tmodels.resnet18(pretrained=True)
|
|
|
+ self.backbone = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
|
|
|
self.backbone.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
|
|
|
self.convolution2d = nn.Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1), bias=True)
|
|
|
self.fc_maxpool = nn.AdaptiveMaxPool2d((1, 1))
|