David310/Detect_AI-generated_Image

1import os2import torch3import torch.nn as nn4from torch.nn import init5from torch.optim import lr_scheduler6 7 8class BaseModel(nn.Module):9    def __init__(self, opt):10        super(BaseModel, self).__init__()11        self.opt = opt12        self.total_steps = 013        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)14        self.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')15 16    def save_networks(self, save_filename):17        save_path = os.path.join(self.save_dir, save_filename)18 19        # serialize model and optimizer to dict20        state_dict = {21            'model': self.model.state_dict(),22            'optimizer' : self.optimizer.state_dict(),23            'total_steps' : self.total_steps,24        }25 26        torch.save(state_dict, save_path)27 28 29    def eval(self):30        self.model.eval()31 32    def test(self):33        with torch.no_grad():34            self.forward()35 36 37def init_weights(net, init_type='normal', gain=0.02):38    def init_func(m):39        classname = m.__class__.__name__40        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):41            if init_type == 'normal':42                init.normal_(m.weight.data, 0.0, gain)43            elif init_type == 'xavier':44                init.xavier_normal_(m.weight.data, gain=gain)45            elif init_type == 'kaiming':46                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')47            elif init_type == 'orthogonal':48                init.orthogonal_(m.weight.data, gain=gain)49            else:50                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)51            if hasattr(m, 'bias') and m.bias is not None:52                init.constant_(m.bias.data, 0.0)53        elif classname.find('BatchNorm2d') != -1:54            init.normal_(m.weight.data, 1.0, gain)55            init.constant_(m.bias.data, 0.0)56 57    print('initialize network with %s' % init_type)58    net.apply(init_func)59