<jupyter_start><jupyter_text>Train Unet on FMD dataset<jupyter_code>%load_ext autoreload
%autoreload 2
import sys
sys.path.append("..")

from util import show, plot_images, plot_tensors
from data_loader_FMD import pil_loader, load_denoising
import torch<jupyter_output><empty_output><jupyter_text>Load the FMD dataset<jupyter_code>root = '/home/yina/DeepDenoising/Data/FMD/'
loader = pil_loader
train = True
noise_levels = [1, 2, 4, 8, 16]
types = None
captures = 20
patch_size = 128
batch_size = 32
transform = None
target_transform = None
test_fov = 19

kwargs = {'drop_last': False}
add_kwargs = {'num_workers': 4, 'pin_memory': True} \
    if torch.cuda.is_available() else {}
kwargs.update(add_kwargs)
loader = load_denoising(root, train, batch_size, noise_levels=noise_levels,
    types=types, captures=captures, patch_size=patch_size, transform=transform, 
    target_transform=target_transform, loader=pil_loader, test_fov=test_fov)
for i_batch, (noisy, clean) in enumerate(loader):
    print(i_batch, noisy.shape)
    print(i_batch, clean.shape)
    plot_tensors([noisy[0,0,:], clean[0,0,:]], ['Noisy Image', 'Clean Image'])
    break<jupyter_output>0 torch.Size([32, 1, 128, 128])
0 torch.Size([32, 1, 128, 128])<jupyter_text>Model<jupyter_code>from models.unet import Unet
model = Unet(up='nearest')<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>from util import getbestgpu
device = getbestgpu(1)
model = model.to(device)
from torch.nn import MSELoss
from torch.optim import Adam

loss_function = MSELoss()
optimizer = Adam(model.parameters(), lr=0.001)
n_epoch = 50
for epoch in range(n_epoch):

    running_loss = 0.0
    
    for i, (noisy, clean) in enumerate(loader):
        noisy_images, clean_images = noisy.type(torch.FloatTensor).to(device), clean.type(torch.FloatTensor).to(device)
        
        optimizer.zero_grad()
    
        net_output = model(noisy_images)
        loss = loss_function(net_output, clean_images)
 
        loss.backward()
    
        optimizer.step()
    
         # print statistics
        running_loss += loss.item()
        if i % 200 == 199:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.5f' %
                  (epoch + 1, i + 1, running_loss / 200))
            print(net_output.shape)
            running_loss = 0.0

print('Finished Training')
plot_tensors([noisy_images[0,0,:], net_output[0,0,:], clean_images[0,0,:]], ['Input','Output','Target'])
torch.save(model.state_dict(), './models/FMD_epoch50_model_nearestup')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1293610348882791, 'fraction_numerical': 0.03018424147393179, 'mean_word_length': 4.4297872340425535, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '564950', 'license': 'MIT', 'n_tokens_mistral': 968, 'n_tokens_neox': 900, 'n_words': 186, 'path': 'Self-Supervised Training with Noise2Self/notebooks/Train-Unet-on-FMD.ipynb', 'repo_name': 'BoHuangLab/timeunet', 'size': '2551'}