Hello, when I add a learning rate scheduler, the error appeared. I guess the part should be fixed are in
save_checkpoint or load_checkpoint shown below. How should I fixed the error? Thanks.
Functions:
import torch
import config
import torch.nn as nn
import random
import os
import numpy as np
from torchvision.utils import save_image
from skimage.metrics import structural_similarity as ssim
def save_checkpoint(model, optimizer, filename="my_checkpoint.pth.tar"):
print("=> Saving checkpoint")
checkpoint = {
"state_dict": model.state_dict(),
"optimizer": optimizer.state_dict(),
}
torch.save(checkpoint, filename)
def load_checkpoint(checkpoint_file, model, optimizer, lr):
print("=> Loading checkpoint")
checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE)
model.load_state_dict(checkpoint["state_dict"])
optimizer.load_state_dict(checkpoint["optimizer"])
# If we don't do this then it will just have learning rate of old checkpoint
# and it will lead to many hours of debugging \:
for param_group in optimizer.param_groups:
param_group["lr"] = lr
Training script:
from model3 import AutoEncoder
import torch
import numpy as np
from dataset import SurDataset
from utils import save_checkpoint, load_checkpoint, save_some_examples, seed_everything, initialize_weights
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
import config
import matplotlib.pyplot as plt
from tqdm import tqdm
from utils import ssim_for_img_tensor
loss_list = []
ssim_list = []
def train_fn(epoch, loader, val_loader, model, optimizer, scheduler, loss_fn, scaler):
loop = tqdm(loader)
for idx, (csv_, target) in enumerate(loop):
csv_ = csv_.to(config.DEVICE)
target = target.to(config.DEVICE)
with torch.cuda.amp.autocast():
model.eval()
predict = model(csv_)
loss = loss_fn(predict, target)
model.train()
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.step(scheduler)
scaler.update()
ssim=ssim_for_img_tensor(target, predict, mode='sigmoid')
loss_list.append(loss.item())
ssim_list.append(ssim)
loop.set_description(f"{epoch+1}/{config.NUM_EPOCHS}")
loop.set_postfix(loss=loss.item(), batch_avg_ssim=ssim)
if idx % 1 == 0:
save_some_examples(model, val_loader, epoch, folder=config.EVALUATE_FOLDERNAME, mode='sigmoid')
if config.SAVE_MODEL:
save_checkpoint(model, optimizer, filename=config.CHECKPOINT)
################################################## plot ###################################################
plt.figure()
plt.subplot(211)
plt.plot(np.arange(1, epoch+2), loss_list, label='MSE Loss')
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.title("Loss vs Epochs Trend")
plt.legend(loc = "best", fontsize=10)
plt.xticks(np.arange(1, config.NUM_EPOCHS+1))
plt.subplot(212)
plt.plot(np.arange(1, epoch+2), ssim_list, label='SSIM Score')
plt.xlabel("Epochs")
plt.ylabel("SSIM Score")
plt.title("SSIM vs Epochs Trend")
plt.legend(loc = "best", fontsize=10)
plt.xticks(np.arange(1, config.NUM_EPOCHS+1))
plt.subplots_adjust(left=0.125,
bottom=0.1,
right=0.9,
top=1,
wspace=0.2,
hspace=0.35)
plt.savefig(f"{config.MODELNAME}_training_process.png")
################################################## plot ###################################################
def main():
seed_everything(42)
print(config.DEVICE)
model = AutoEncoder(in_channels=1, out_channels=1).to(config.DEVICE)
initialize_weights(model)
optimizer = optim.Adam(list(model.parameters()),
lr=config.LEARNING_RATE,
betas=(0.5, 0.999))
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[7, 15, 23], gamma=0.1)
mse = nn.MSELoss()
scaler = torch.cuda.amp.GradScaler()
dataset = SurDataset(csv_dir=config.CSV_DIR, img_dir=config.IMG_DIR, rescale=False)
train_set, val_set = torch.utils.data.random_split(dataset, [12000, len(dataset) - 12000])
loader = DataLoader(dataset=train_set, batch_size=config.BATCH_SIZE, shuffle=True)
val_loader = DataLoader(dataset=val_set, batch_size=config.BATCH_SIZE)
if config.LOAD_MODEL:
load_checkpoint(
config.CHECKPOINT, model, optimizer, config.LEARNING_RATE,
)
model.train()
for epoch in range(config.NUM_EPOCHS):
train_fn(epoch=epoch, loader=loader, val_loader=val_loader, model=model, optimizer=optimizer, scheduler=scheduler, loss_fn=mse, scaler=scaler)
if __name__=='__main__':
main()
Errors:
Traceback (most recent call last):
File "C:\Users\PML\.conda\envs\floren\lib\runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "C:\Users\PML\.conda\envs\floren\lib\runpy.py", line 85, in _run_code
exec(code, run_globals)
File "c:\Users\PML\.vscode\extensions\ms-python.python-2021.6.944021595\pythonFiles\lib\python\debugpy\__main__.py", line 45, in <module>
cli.main()
File "c:\Users\PML\.vscode\extensions\ms-python.python-2021.6.944021595\pythonFiles\lib\python\debugpy/..\debugpy\server\cli.py", line 444, in main
run()
File "c:\Users\PML\.vscode\extensions\ms-python.python-2021.6.944021595\pythonFiles\lib\python\debugpy/..\debugpy\server\cli.py", line 285, in run_file
runpy.run_path(target_as_str, run_name=compat.force_str("__main__"))
File "C:\Users\PML\.conda\envs\floren\lib\runpy.py", line 263, in run_path
pkg_name=pkg_name, script_name=fname)
File "C:\Users\PML\.conda\envs\floren\lib\runpy.py", line 96, in _run_module_code
mod_name, mod_spec, pkg_name, script_name)
File "C:\Users\PML\.conda\envs\floren\lib\runpy.py", line 85, in _run_code
exec(code, run_globals)
File "c:\Users\PML\Documents\Florentino_space\cast_iron_preprocess\autoencoder20210615\train2.py", line 96, in <module>
main()
File "c:\Users\PML\Documents\Florentino_space\cast_iron_preprocess\autoencoder20210615\train2.py", line 93, in main
train_fn(epoch=epoch, loader=loader, val_loader=val_loader, model=model, optimizer=optimizer, scheduler=scheduler, loss_fn=mse, scaler=scaler)
File "c:\Users\PML\Documents\Florentino_space\cast_iron_preprocess\autoencoder20210615\train2.py", line 30, in train_fn
scaler.step(scheduler)
File "C:\Users\PML\.conda\envs\floren\lib\site-packages\torch\cuda\amp\grad_scaler.py", line 316, in step
self.unscale_(optimizer)
File "C:\Users\PML\.conda\envs\floren\lib\site-packages\torch\cuda\amp\grad_scaler.py", line 267, in unscale_
optimizer_state["found_inf_per_device"] = self._unscale_grads_(optimizer, inv_scale, found_inf, False)
File "C:\Users\PML\.conda\envs\floren\lib\site-packages\torch\cuda\amp\grad_scaler.py", line 194, in _unscale_grads_
for group in optimizer.param_groups:
AttributeError: 'MultiStepLR' object has no attribute 'param_groups'