当前位置：首页 > news >正文

SCINet 训练代码修改

news 来源：原创 2025/5/7 12:59:16

不多说，放代码

import os
import sys
import time
import glob
import numpy as np
import torch
import utils
from PIL import Image
import logging
import argparse
import torch.utils
import torch.backends.cudnn as cudnn
import torch.nn as nn
from torch.autograd import Variablefrom model import *
from multi_read_data import MemoryFriendlyLoaderparser = argparse.ArgumentParser("SCI")
parser.add_argument('--batch_size', type=int, default=16, help='batch size')
parser.add_argument('--cuda', default=True, type=bool, help='Use CUDA to train model')
parser.add_argument('--gpu', type=str, default='0', help='gpu device id')
parser.add_argument('--seed', type=int, default=2, help='random seed')
parser.add_argument('--epochs', type=int, default=100, help='epochs')
parser.add_argument('--lr', type=float, default=0.0003, help='learning rate')
parser.add_argument('--stage', type=int, default=3, help='epochs')
parser.add_argument('--save', type=str, default='EXP/', help='location of the data corpus')args = parser.parse_args()os.environ["CUDA_VISIBLE_DEVICES"] = args.gpuargs.save = args.save + '/' + 'Train-{}'.format(time.strftime("%Y%m%d-%H%M%S"))
utils.create_exp_dir(args.save, scripts_to_save=glob.glob('*.py'))
model_path = args.save + '/model_epochs/'
os.makedirs(model_path, exist_ok=True)
image_path = args.save + '/image_epochs/'
os.makedirs(image_path, exist_ok=True)log_format = '%(asctime)s %(message)s'
logging.basicConfig(stream=sys.stdout, level=logging.INFO,format=log_format, datefmt='%m/%d %I:%M:%S %p')
fh = logging.FileHandler(os.path.join(args.save, 'log.txt'))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)logging.info("train file name = %s", os.path.split(__file__))if torch.cuda.is_available():if args.cuda:torch.set_default_tensor_type('torch.cuda.FloatTensor')if not args.cuda:print("WARNING: It looks like you have a CUDA device, but aren't " +"using CUDA.\nRun with --cuda for optimal training speed.")torch.set_default_tensor_type('torch.FloatTensor')
else:torch.set_default_tensor_type('torch.FloatTensor')def save_images(tensor, path):image_numpy = tensor[0].cpu().float().numpy()image_numpy = (np.transpose(image_numpy, (1, 2, 0)))im = Image.fromarray(np.clip(image_numpy * 255.0, 0, 255.0).astype('uint8'))im.save(path, 'png')def main():if not torch.cuda.is_available():logging.info('no gpu device available')sys.exit(1)np.random.seed(args.seed)cudnn.benchmark = Truetorch.manual_seed(args.seed)cudnn.enabled = Truetorch.cuda.manual_seed(args.seed)logging.info('gpu device = %s' % args.gpu)logging.info("args = %s", args)model = Network(stage=args.stage)model.enhance.in_conv.apply(model.weights_init)model.enhance.conv.apply(model.weights_init)model.enhance.out_conv.apply(model.weights_init)model.calibrate.in_conv.apply(model.weights_init)model.calibrate.convs.apply(model.weights_init)model.calibrate.out_conv.apply(model.weights_init)model = model.cuda()optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999), weight_decay=3e-4)MB = utils.count_parameters_in_MB(model)logging.info("model size = %f", MB)print(MB)train_low_data_names = '/root/autodl-tmp/our485/low'TrainDataset = MemoryFriendlyLoader(img_dir=train_low_data_names, task='train')test_low_data_names = '/root/autodl-tmp/eval15/low'TestDataset = MemoryFriendlyLoader(img_dir=test_low_data_names, task='test')# 创建 CUDA 随机数生成器g = torch.Generator(device='cuda')g.manual_seed(args.seed)train_queue = torch.utils.data.DataLoader(TrainDataset, batch_size=args.batch_size,pin_memory=True, num_workers=0, shuffle=True, generator=g)test_queue = torch.utils.data.DataLoader(TestDataset, batch_size=1,pin_memory=True, num_workers=0, shuffle=True, generator=g)total_step = 0for epoch in range(args.epochs):model.train()losses = []for batch_idx, (input, _) in enumerate(train_queue):total_step += 1input = Variable(input, requires_grad=False).cuda()optimizer.zero_grad()loss = model._loss(input)loss.backward()nn.utils.clip_grad_norm_(model.parameters(), 5)optimizer.step()losses.append(loss.item())logging.info('train-epoch %03d %03d %f', epoch, batch_idx, loss)logging.info('train-epoch %03d %f', epoch, np.average(losses))utils.save(model, os.path.join(model_path, 'weights_%d.pt' % epoch))if epoch % 50 == 0 and total_step != 0:logging.info('train %03d %f', epoch, loss)model.eval()with torch.no_grad():for _, (input, image_name) in enumerate(test_queue):input = Variable(input, volatile=True).cuda()# image_name = image_name[0].split('\\')[-1].split('.')[0]image_name_str = image_name[0]# 使用上述方法处理 image_name_strimage_name = os.path.basename(image_name_str)  # 这里使用 os.path.basename 方法image_name = image_name.split('.')[0]  # 如果还需要去掉文件扩展名，可以再进行一次分割illu_list, ref_list, input_list, atten = model(input)u_name = '%s.png' % (image_name + '_' + str(epoch))u_path = os.path.join(image_path, u_name)  save_images(ref_list[0], u_path)if __name__ == '__main__':main()