Thanks. Please see below. I hope this is reproduceable:
import argparse
import re
import os, glob, datetime, time
os.environ[‘KMP_DUPLICATE_LIB_OK’]=‘True’
os.environ[‘KMP_DUPLICATE_LIB_OK’]=‘True’
os.environ[“CUDA_VISIBLE_DEVICES”] = “0”
import numpy as np
import torch
import torch.nn as nn
from torch.nn.modules.loss import _Loss
import torch.nn.init as init
from torch.utils.data import DataLoader
import torch.optim as optim
from torch.optim.lr_scheduler import MultiStepLR
import multiprocessing
import scipy.io
from torch.utils.data import Dataset
class Normalization:
def init(self, mean_val=None,std_val=None):
self.mean_val = mean_val
self.std_val = std_val
def normalize(self, x):
return (x-self.mean_val)/self.std_val
def unnormalize(self, x):
return x*self.std_val + self.mean_val
parser = argparse.ArgumentParser(description=‘PyTorch DnCNN’)
parser.add_argument(‘–model’, default=‘DnCNN’, type=str, help=‘choose a type of model’)
parser.add_argument(‘–data_dir’, default=‘data/train’, type=str, help=‘path of train data’)
parser.add_argument(‘–sigma’, default=25, type=int, help=‘noise level’)
parser.add_argument(‘–sigma2’, default=35, type=int, help=‘noise level2’)
parser.add_argument(‘–epoch’, default=30, type=int, help=‘number of train epoches’)
parser.add_argument(‘–lr’, default=1e-3, type=float, help=‘initial learning rate for Adam’)
parser.add_argument(‘–batch_size’, default=5, type=int, help=‘batch size’)
parser.add_argument(‘–patch_size’, default=(3000,240), type=int, help=‘patch size’)
parser.add_argument(‘–stride’, default=(32,32), type=int, help=‘the step size to slide on the data’)
parser.add_argument(‘–jump’, default=3, type=int, help=‘the space between shot’)
parser.add_argument(‘–download’, default=False, type=bool, help=‘if you will download the dataset from the internet’)
parser.add_argument(‘–datasets’, default = 0, type = int, help=‘the num of datasets you want be download,if download = True’)
parser.add_argument(‘–train_data_num’, default=4000, type=int, help=‘the num of the train_data’)
parser.add_argument(‘–aug_times’, default=0, type=int, help=‘Number of aug operations’)
parser.add_argument(‘–scales’, default=[1], type=list, help=‘data scaling’)
parser.add_argument(‘–agc’, default=True, type=int, help=‘Normalize each trace by amplitude’)
parser.add_argument(‘–verbose’, default=True, type=int, help=‘Whether to output the progress of data generation’)
parser.add_argument(‘–display’, default=5, type=int, help=‘interval for displaying loss’)
args = parser.parse_args()
batch_size = args.batch_size
cuda = torch.cuda.is_available()
torch.set_default_dtype(torch.float64)
n_epoch = args.epoch
sigma = args.sigma
sigma2 = args.sigma2
if not os.path.exists(‘models_denoise’):
os.mkdir(‘models_denoise’)
save_dir = os.path.join(‘models_denoise’, args.model+‘_’ + ‘noNoise’ )
if not os.path.exists(save_dir):
os.mkdir(save_dir)
class DenoisingDataset(Dataset):
“”“Dataset wrapping tensors.
Arguments:
xs (Tensor): clean data patches
sigma: noise level, e.g., 25
“””
def init(self, xs, xs2):
super(DenoisingDataset, self).init()
self.xs = xs
self.xs2 = xs2
def __getitem__(self, index):
batch_x = self.xs[index]
batch_y=self.xs2[index]
return batch_x, batch_y
def __len__(self):
return self.xs.size(0)
class DnCNN(nn.Module):
def init(self, depth=17, n_channels=64, image_channels=1, use_bnorm=True, kernel_size=3):
super(DnCNN, self).init()
kernel_size = 3
padding = 1
layers =
layers.append(nn.Conv2d(in_channels=image_channels, out_channels=n_channels, kernel_size=kernel_size, padding=padding, bias=True))
layers.append(nn.ReLU(inplace=True))
for _ in range(depth-2):
layers.append(nn.Conv2d(in_channels=n_channels, out_channels=n_channels, kernel_size=kernel_size, padding=padding, bias=False))
layers.append(nn.BatchNorm2d(n_channels, eps=0.0001, momentum = 0.95))
layers.append(nn.ReLU(inplace=True))
layers.append(nn.Conv2d(in_channels=n_channels, out_channels=image_channels, kernel_size=kernel_size, padding=padding, bias=False))
self.dncnn = nn.Sequential(*layers)
self._initialize_weights()
def forward(self, x):
y = x
out = self.dncnn(x)
return y-out
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.orthogonal_(m.weight)
if m.bias is not None:
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
print('init weight')
def findLastCheckpoint(save_dir):
file_list = glob.glob(os.path.join(save_dir, 'model_.pth’))
if file_list:
epochs_exist = []
for file_ in file_list:
result = re.findall(".model_(.).pth.", file_)
epochs_exist.append(int(result[0]))
initial_epoch = max(epochs_exist)
else:
initial_epoch = 0
return initial_epoch
def log(*args, **kwargs):
print(datetime.datetime.now().strftime(“%Y-%m-%d %H:%M:%S:”), *args, **kwargs)
if name == ‘main’:
print(‘===> Building model’)
device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
model = DnCNN()
initial_epoch = findLastCheckpoint(save_dir=save_dir) # load the last model in matconvnet style
if initial_epoch > 1:
print('resuming by loading epoch %03d\n' % (initial_epoch-1))
if initial_epoch >= n_epoch:
print("training have finished")
else:
model = torch.load(os.path.join(save_dir, 'model_%03d.pth' % (initial_epoch-1)))
model.train()
criterion = nn.MSELoss(reduce = True,size_average = False)
if cuda:
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2) # learning rates
xss=np.random.random((400,3000,240)).astype(np.float32)
xss2=np.random.random((400,3000,240)).astype(np.float32)
xs=xss.transpose(0,2,1)
xs_mean= torch.tensor(np.mean(xs,axis=(0,-1),keepdims=True)).to(device)
xs_std=torch.tensor(np.std(xs,axis=(0,-1),keepdims=True)).to(device)
xs_normalization = Normalization(mean_val=xs_mean,std_val=xs_std)
xs=torch.tensor(xs).to(device)
xs=xs_normalization.normalize(xs)
xs=torch.unsqueeze(xs,1)
xs2=xss2.transpose(0,2,1)
xs2_mean= torch.tensor(np.mean(xs2,axis=(0,-1),keepdims=True)).to(device)
xs2_std=torch.tensor(np.std(xs2,axis=(0,-1),keepdims=True)).to(device)
xs2_normalization = Normalization(mean_val=xs2_mean,std_val=xs2_std)
xs2=torch.tensor(xs2).to(device)
xs2=xs2_normalization.normalize(xs2)
xs2=torch.unsqueeze(xs2,1)
for epoch in range(initial_epoch, n_epoch):
scheduler.step(epoch) # step to the learning rate in this epcoh
DDataset = DenoisingDataset(xs, xs2)
DLoader = DataLoader(dataset=DDataset, drop_last=True, batch_size=batch_size, shuffle=True)
epoch_loss = 0
start_time = time.time()
for n_count, batch_yx in enumerate(DLoader):
optimizer.zero_grad()
if cuda:
batch_y, batch_x = batch_yx[0].cuda(), batch_yx[1].cuda()
else:
batch_x, batch_y = batch_yx[0], batch_yx[1] # #batchx is odd batcg y even
loss = criterion(model(batch_x), batch_y)
epoch_loss += loss.item()
loss.backward()
optimizer.step()