def __init__(self, backbone='vgg_bn'):
super().__init__()
self.backbone_name = backbone
if backbone == "vgg" or backbone == 'vgg_bn':
if backbone == 'vgg_bn':
self.backbone = VggNet(name="vgg16_bn", pretrain=True)
elif backbone == 'vgg':
self.backbone = VggNet(name="vgg16", pretrain=True)
self.deconv5 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1)
self.merge4 = UpBlok(512 + 256, 128)
self.merge3 = UpBlok(256 + 128, 64)
self.merge2 = UpBlok(128 + 64, 32)
self.merge1 = UpBlok(64 + 32, 16)
elif backbone == 'resnet50' or backbone == 'resnet101':
if backbone == 'resnet101':
self.backbone = ResNet(name="resnet101", pretrain=True)
elif backbone == 'resnet50':
self.backbone = ResNet(name="resnet50", pretrain=True)
self.deconv5 = nn.ConvTranspose2d(2048, 256, kernel_size=4, stride=2, padding=1)
self.merge4 = UpBlok(1024 + 256, 128)
self.merge3 = UpBlok(512 + 128, 64)
self.merge2 = UpBlok(256 + 64, 32)
self.merge1 = UpBlok(64 + 32, 16)
else:
print("backbone is not support !")
def se_resnet152(num_classes=1_000):
"""Constructs a ResNet-152 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(SEBottleneck, [3, 8, 36, 3], num_classes=num_classes)
model.avgpool = nn.AdaptiveAvgPool2d(1)
return model
def se_resnet34(num_classes=1_000):
"""Constructs a ResNet-34 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(SEBasicBlock, [3, 4, 6, 3], num_classes=num_classes)
model.avgpool = nn.AdaptiveAvgPool2d(1)
return model
def se_resnet50(num_classes=110, pretrained=False):
"""Constructs a ResNet-50 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(SEBottleneck, [3, 4, 6, 3], num_classes=num_classes)
model.avgpool = nn.AdaptiveAvgPool2d(1)
if pretrained:
model.load_state_dict(
model_zoo.load_url(
"https://www.dropbox.com/s/xpq8ne7rwa4kg4c/seresnet50-60a8950a85b2b.pkl"
))
return model
def main():
# import arguments for experiment
parser = BYOL_parser_COVID()
args = parser.parse_args()
# Hold randomness for reproducible experiment
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
# torch.cuda.manual_seed_all(args.seed) # if use multi-GPU
cudnn.deterministic = True
cudnn.benchmark = False
np.random.seed(args.seed) # numpy seed affects torchvision.transforms
random.seed(args.seed) # random seed affects data loader's sampling
device = args.device
print(f"Training with: {device}")
# set encoder for online network & target network
encoder = ResNet(args)
# load pre-trained model if defined :TODO
# set predictor network
predictor = MLP_Head(in_channels=args.mlp_projection_size,
mlp_hidden_size=args.mlp_hidden_size,
projection_size=args.mlp_projection_size)
# set optimizer
optimizer = SGD
# set learner
trainer = BYOL(args, encoder, predictor, optimizer)
# train learner
trainer.train()
def train():
""" train """
print('execute train')
# TODO
train_inputs = None
train_teachers = None
test_inputs = None
test_teachers = None
network = ResNet(INPUT_SHAPE)
network.print_model_summay()
model = network.get_model()
model.compile(optimizer=SGD(momentum=0.9, decay=0.0001),
loss='categorical_crossentropy',
metrics=['accuracy'])
plot_model(model, to_file=os.path.join(os.pardir, 'ModelLayers.png'))