def _init_model(self):
criterion = nn.BCELoss()
self.criterion = criterion.to(self.device)
M = Models()
model = M.PSP(img_ch=3, output_ch=1)
self.model = model.to(self.device)
# init_weights(self.model, 'kaiming', gain=1)
# summary(self.model, input_size=(4, 448, 448))
self.model_optimizer = optim.Adamax(model.parameters(),
lr=1e-3,
weight_decay=0.01)
self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
self.model_optimizer, T_max=len(self.train_queue))
