class ResNet50_FasterRCNN:
def __init__(self, pretrained=False):
# Building our FasterRCNN model for objects detection
backbone = resnet_fpn_backbone('resnet50', pretrained=pretrained)
num_classes = 4 + 1
anchor_generator = AnchorGenerator(sizes=(40, 60, 150, 200, 250),
aspect_ratios=(0.7, 1.0, 1.3))
self.model = FRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator)
def train(self):
self.model.train()
def to(self, device):
self.model.to(device)
def eval(self):
self.model.eval()
def parameters(self):
return self.model.parameters()
def get_state_dict(self):
return self.model.state_dict()
def set_state_dict(self, state_dict):
self.model.load_state_dict(state_dict)
def fit_batch(self, images, target):
return self.model(images, target)
def predict_batch(self, images):
return self.model(images)
def main(network):
# train on the GPU or on the CPU, if a GPU is not available
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
# our dataset has two classes only - background and person
num_classes = 2
#dataset = torch.utils.data.Subset(TBdata,[range(len(TBdata))])
indices = torch.randperm(len(TBdata)).tolist()
dataset = torch.utils.data.Subset(TBdata, indices[:])
indices_ = torch.randperm(len(TBdata_test)).tolist()
dataset_val = torch.utils.data.Subset(TBdata_test, indices_[:])
# get the model using our helper function
#model = get_model_instance_segmentation(num_classes)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=8,
sampler=None,
num_workers=0,
collate_fn=collate_fn)
dataloader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=8,
sampler=None,
num_workers=0,
collate_fn=collate_fn)
#Calculated statistics on training data:
#Transform parameters
min_size = 550
max_size = 700
image_means = [0.9492, 0.9492, 0.9492]
image_stds = [0.1158, 0.1158, 0.1158]
if network == 'resnet50':
backbone = resnet_fpn_backbone('resnet50', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'resnet18':
backbone = resnet_fpn_backbone('resnet18', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'resnet152':
backbone = resnet_fpn_backbone('resnet152', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'RPNresnet50':
backbone = resnet_fpn_backbone('resnet50', True)
model = RPN_custom(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'RPNresnet152':
backbone = resnet_fpn_backbone('resnet152', True)
model = RPN_custom(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
# move model to the right device
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
num_epochs = 10
Ls = {
'total loss': [],
'loss_classifier': [],
'loss_box_reg': [],
'loss_objectness': [],
'loss_rpn_box_reg': []
}
Ls_val = {
'total loss': [],
'loss_classifier': [],
'loss_box_reg': [],
'loss_objectness': [],
'loss_rpn_box_reg': []
}
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model,
optimizer,
dataloader,
device,
epoch,
print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
#evaluate(model, dataloader_test, device=device)
Ls_val = record_losses(model, dataloader_val, device, Ls_val, network)
#record losses
Ls = record_losses(model, dataloader, device, Ls, network)
#If folder does not exist already, create it
output_loc = f'./{network}/'
if not os.path.exists(output_loc):
os.makedirs(output_loc)
torch.save(model.state_dict(), output_loc + 'model.pt')
print("That's it!")
return Ls, Ls_val, num_epochs
class FasterRCNNFood:
def __init__(self,
backbone_name: str,
pretrained: bool = True,
finetune: bool = True,
num_classes: int = 2):
self.__pretrained = pretrained
self.__num_classes = num_classes
self.__model_name = backbone_name
backbone = build_backbone(backbone_name, pretrained, finetune)
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
self.model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
self.params = [p for p in self.model.parameters() if p.requires_grad]
self.optimizer = torch.optim.Adam(params=self.params,
lr=0.005,
weight_decay=0.0005)
self.lr_scheduler = torch.optim.lr_scheduler.StepLR(
optimizer=self.optimizer, step_size=3, gamma=0.1)
def train(self,
data_loader: DataLoader,
data_loader_test: DataLoader,
num_epochs: int = 10,
use_cuda: bool = True,
epoch_save_ckpt: Union[int, list] = None,
dir: str = None):
"""
Method to train FasterRCNNFood model.
Args:
data_loader (torch.utils.data.DataLoader): data loader to train model on
data_loader_test (torch.utils.data.DataLoader): data loader to evaluate model on
num_epochs (int = 10): number of epoch to train model
use_cuda (bool = True): use cuda or not
epoch_save_ckpt (list or int): Epoch at which you want to save the model. If -1 save only last epoch.
dir (str = "models/): Directory where model are saved under the name "{model_name}_{date}_ep{epoch}.pth"
"""
if epoch_save_ckpt == -1:
epoch_save_ckpt = [num_epochs - 1]
if not dir:
dir = "models"
dir = Path(dir)
dir.mkdir(parents=True, exist_ok=True)
# choose device
if use_cuda and torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
# define dataset
self.model.to(device)
writer = SummaryWriter()
for epoch in range(num_epochs):
# train for one epoch, printing every 50 iterations
train_one_epoch(self.model,
self.optimizer,
data_loader,
device,
epoch,
print_freq=50,
writer=writer)
# update the learning rate
self.lr_scheduler.step()
# evaluate on the test dataset
evaluate(self.model,
data_loader_test,
device=device,
writer=writer,
epoch=epoch)
# save checkpoint
if epoch in epoch_save_ckpt:
self.save_checkpoint(dir.as_posix(), epoch)
writer.close()
print("That's it!")
def save_checkpoint(self, dir: str, epoch: int):
"""
Save a model checkpoint at a given epoch.
Args:
dir: dir folder to save the .pth file
epoch: epoch the model is
"""
state = {
'epoch': epoch + 1,
'state_dict': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'num_classes': self.__num_classes,
'pretrained': self.__pretrained,
"model_name": self.__model_name
}
now = datetime.now()
filename = "{model_name}_{date}_ep{epoch}.pth".format(
model_name=self.__model_name,
date=now.strftime("%b%d_%H-%M"),
epoch=epoch)
torch.save(state, Path(dir) / filename)
"Checkpoint saved : {}".format(Path(dir) / filename)
def predict(self, dataset, idx):
img, _ = dataset[idx]
img.to("cpu")
self.model.eval()
self.model.to("cpu")
pred = self.model([img])
return img, pred[0]
@staticmethod
def load_checkpoint(filename: str,
cuda: bool = True) -> ("FasterRCNNFood", int):
"""
Load a model checkpoint to continue training.
Args:
filename (str): filename/path of the checkpoint.pth
cuda (bool = True): use cuda
Returns:
(FasterRCNNFood) model
(int) number of epoch + 1 the model was trained with
"""
device = torch.device("cuda") if (
cuda and torch.cuda.is_available()) else torch.device("cpu")
start_epoch = 0
if Path(filename).exists():
print("=> loading checkpoint '{}'".format(filename))
checkpoint = torch.load(filename, map_location=device)
# Load params
pretrained = checkpoint['pretrained']
num_classes = checkpoint["num_classes"]
start_epoch = checkpoint['epoch']
model_name = checkpoint['model_name']
# Build model key/architecture
model = FasterRCNNFood(model_name, pretrained, num_classes)
# Update model and optimizer
model.model.load_state_dict(checkpoint['state_dict'])
model.optimizer.load_state_dict(checkpoint['optimizer'])
model.model = model.model.to(device)
# now individually transfer the optimizer parts...
for state in model.optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.to(device)
print("=> loaded checkpoint '{}' (epoch {})".format(
filename, checkpoint['epoch']))
return model, start_epoch
else:
print("=> no checkpoint found at '{}'".format(filename))
@staticmethod
def load_for_inference(filename: str,
cuda: bool = True) -> "FasterRCNNFood":
"""
Load a model checkpoint to make inference.
Args:
filename (str): filename/path of the checkpoint.pth
cuda (bool = True): use cuda
Returns:
(FasterRCNNFood) model
"""
device = torch.device("cuda") if (
cuda and torch.cuda.is_available()) else torch.device("cpu")
if Path(filename).exists():
print("=> loading checkpoint '{}'".format(filename))
checkpoint = torch.load(filename, map_location=device)
# Load params
pretrained = checkpoint['pretrained']
num_classes = checkpoint["num_classes"]
model_name = checkpoint['model_name']
# Build model key/architecture
model = FasterRCNNFood(model_name, pretrained, num_classes)
# Update model and optimizer
model.model.load_state_dict(checkpoint['state_dict'])
model.model = model.model.to(device)
model.model = model.model.eval()
print("=> loaded checkpoint '{}'".format(filename))
return model
else:
print("=> no checkpoint found at '{}'".format(filename))
backbone = backboneNet_efficient() # use efficientnet as our backbone
backboneFPN = backboneWithFPN(backbone) # add FPN
anchor_generator = AnchorGenerator(cfg.anchor_sizes, cfg.aspect_ratios)
model_ft = FasterRCNN(backboneFPN,
num_classes=cfg.num_classes,
rpn_anchor_generator=anchor_generator,
min_size=cfg.min_size,
max_size=cfg.max_size)
model_ft.to(device)
optimizer_ft = optim.SGD(model_ft.parameters(),
lr=cfg.learning_rate,
momentum=cfg.momentum,
weight_decay=cfg.weight_decay)
lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft,
milestones=cfg.milestones,
gamma=cfg.gamma)
model_ft = train_model(model_ft,
train_loader,
valid_loader,
optimizer_ft,
lr_scheduler,
num_epochs=cfg.epochs)
torch.save(model_ft.state_dict(), cfg.model_folder + cfg.model_name)
def get_fasterrcnn_model(arch_str,
num_classes,
pretrained=True,
pretrained_backbone=True,
trainable_layers=5,
**kwargs):
"""Creates FasterRCNN model with resnet backbone"""
#if pretrained == True: pretrained_backbone=False
backbone = resnet_fpn_backbone(arch_str,
pretrained=pretrained_backbone,
trainable_layers=trainable_layers)
anchor_sizes = (
(16, ),
(32, ),
(64, ),
(128, ),
(256, ),
)
aspect_ratios = ((0.5, 1.0, 2.0), ) * len(anchor_sizes)
anchor_generator = AnchorGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
model = FasterRCNN(
backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_fg_iou_thresh=0.5,
box_bg_iou_thresh=0.5,
image_mean=[0.0, 0.0, 0.0], # already normalized by fastai
image_std=[1.0, 1.0, 1.0],
#min_size = 1,
#box_score_thresh=0.6,
**kwargs)
if pretrained:
try:
pretrained_dict = load_state_dict_from_url(
_model_urls['fasterrcnn_' + arch_str + '_fpn_coco'],
progress=True)
model_dict = model.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if (k in model_dict) and (
model_dict[k].shape == pretrained_dict[k].shape)
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
#overwrite_eps(model, 0.0)
for module in model.modules():
if isinstance(module, FrozenBatchNorm2d):
module.eps = 0.0
except Exception as e:
#print(e)
print("No pretrained coco model found for fasterrcnn_" + arch_str)
print("This does not affect the backbone.")
return model.train()
def save_parameters(save_dir: str, model: FasterRCNN, epoch: int):
if not os.path.exists(save_dir):
os.mkdir(save_dir)
save_path = os.path.join(save_dir, 'epoch_{}.pt'.format(epoch))
torch.save(model.state_dict(), save_path)
print('Saved model at path {}'.format(save_path))
if iou(t_box[0], o_box) > IOU:
if t_label == o_label:
if t_label == 1:
tpp += 1
elif t_label == 2:
tpn += 1
else:
if t_label == 1:
fpn += 1
fnp += 1
elif t_label == 2:
fpp += 1
fnn += 1
else:
if t_label == 1:
fnp += 1
elif t_label == 2:
fnn += 1
else:
if t_label == 1:
fnp += 1
elif t_label == 2:
fnn += 1
print(
"Epoch {}\t Pos Precisison:{:.3f}\t Pos Recall:{:.3f}\t Neg Precisison:{:.3f}\t Neg Recall:{:.3f}"
.format(epoch + 1, tpp / (tpp + fpp + 1e-5),
tpp / (tpp + fnp + 1e-5), tpn / (tpn + fpn + 1e-5),
tpn / (tpn + fnn + 1e-5)))
torch.save(model.state_dict(),
'./models/' + 'model_' + str(epoch + 1) + '.pth')
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
# let's train it for 10 epochs
num_epochs = 10
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model,
optimizer,
train_loader,
device,
epoch,
print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, test_loader, device=device)
torch.save(model.state_dict(), "efficient_rcnn_" + str(epoch) + ".pth")
root_url = '/home/dung/DocData/cp/145'
dataset = DocDataset(root_url)
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=1, shuffle=True, num_workers=0)
# criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
for i in range(1000):
print('Epoch {}\n'.format(i))
for j, (images, targets) in enumerate(data_loader):
images = images.to(device)
a = {}
a['boxes'] = targets['boxes'][0].to(device)
a['labels'] = targets['labels'][0].to(device)
output = model(images, [a])
losses = sum(loss for loss in output.values())
if j % 30 == 0:
print('Step {} -- loss_classifier = {} -- loss_box_reg = {} -- loss_objectness = {} -- loss_rpn_box_reg = {}\n'.format(j,
output['loss_classifier'].item(), output['loss_box_reg'].item(), output['loss_objectness'].item(), output['loss_rpn_box_reg'].item()))
optimizer.zero_grad()
losses.backward()
optimizer.step()
if i % 20 == 0:
print('save model')
torch.save(model.state_dict(), '3.pth')
print('done')
test_losses = []
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
# train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=1)
_, loss = train(model, optimizer, data_loader, device, epoch, print_freq=1)
train_losses.append(loss)
# # update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, data_loader_validation, device=device)
if (epoch + 1) % 5 == 0:
visualize(model, data_loader_validation, device, epoch,
save_folder=os.path.join(root_path, 'prediction_visualization'))
import matplotlib.pyplot as plt
# Plot training and val loss as a function of the epoch. Use this to monitor for overfitting.
plt.plot(range(1, num_epochs + 1), train_losses, label='traning')
plt.legend(loc='best')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss')
plt.savefig(os.path.join(root_path, 'prediction_visualization', 'loss.png'), bbox_inches='tight')
torch.save(model.state_dict(), "model_mixed_epoch{}.pt".format(20))
# test on testset
evaluate(model, data_loader_test, device=device)
model.train()
loss_dict = model(images, targets)
losses = sum(loss for loss in loss_dict.values())
print("===Epoch:{}/{}===Step:{}/{}===Loss:{:.4f}".format(
ep, epoch, idi, len(dataloader), losses.item()))
# print(loss_dict)
optimizer.zero_grad()
losses.backward()
optimizer.step()
# model.eval()
# prediction = model(images)
# print(prediction[0])
torch.save(model.state_dict(), "./weight_2/{}.pt".format(ep))
# img_array = predictions.permute(1,2,0).detach().cpu().numpy().astype(uint8)
# cv2.imshow("img", cv2.fromarray(img_array))
if cfg.predict:
img_path = "../../data/SUMIT/rs_images_sampled/"
dataset = os.listdir(img_path)
indices = torch.randperm(len(dataset)).tolist()
model.load_state_dict(torch.load("./weight_2/9.pt"), strict=False)
model.to(device)
model.eval()
for idi in indices[-10:]:
img = Image.open(img_path + dataset[idi]).convert("RGB")
# print(np.array(img).shape)
def main():
parser = argparse.ArgumentParser(
description='VISUM 2019 competition - baseline training script',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-d',
'--data_path',
default='/home/master/dataset/train',
metavar='',
help='data directory path')
parser.add_argument('-m',
'--model_path',
default='./model.pth',
metavar='',
help='model file (output of training)')
parser.add_argument('--epochs',
default=60,
type=int,
metavar='',
help='number of epochs') #50
parser.add_argument('--lr',
default=0.01,
type=float,
metavar='',
help='learning rate') #0.005
parser.add_argument('--l2',
default=0.0005,
type=float,
metavar='',
help='L-2 regularization')
args = vars(parser.parse_args())
# Data augmentation
def get_transform(train):
transforms = []
# converts the image, a PIL image, into a PyTorch Tensor
transforms.append(T.ToTensor())
if train:
# during training, randomly flip the training images
# and ground-truth for data augmentation
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)
# backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# backbone.out_channels = 1280
# anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
# aspect_ratios=((0.5, 1.0, 2.0),))
# roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
# output_size=7,
# sampling_ratio=2)
# # put the pieces together inside a FasterRCNN model
# model = FasterRCNN(backbone,
# num_classes=11,
# rpn_anchor_generator=anchor_generator,
# box_roi_pool=roi_pooler)
backbone = torchvision.models.detection.backbone_utils.resnet_fpn_backbone(
'resnet50', True)
backbone.out_channels = 256
anchor_generator = AnchorGenerator(sizes=(8, 16, 32, 64, 128),
aspect_ratios=(0.5, 1.0, 2.0))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# put the pieces together inside a FasterRCNN model
model = FasterRCNN(backbone,
num_classes=11,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# See the model architecture
print(model)
# use our dataset and defined transformations
dataset = VisumData(args['data_path'],
modality='all',
transforms=get_transform(train=True))
dataset_val = VisumData(args['data_path'],
modality='all',
transforms=get_transform(train=False))
# split the dataset in train and test set
torch.manual_seed(1)
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-100])
dataset_val = torch.utils.data.Subset(dataset_val, indices[-100:])
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=2,
shuffle=True,
num_workers=0,
collate_fn=utils.collate_fn)
data_loader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=2,
shuffle=False,
num_workers=0,
collate_fn=utils.collate_fn)
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=args['lr'],
momentum=0.9,
weight_decay=args['l2'])
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=7,
gamma=0.5)
for epoch in range(args['epochs']):
# train for one epoch, printing every 10 iterations
epoch_loss = train_one_epoch(model,
optimizer,
data_loader,
device,
epoch,
print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluator = evaluate(model, data_loader_val, device=device)
# save the model
torch.save(model.state_dict(), args['model_path'])
# cv2.putText(truth, str(labels[i]), (int(x0), int(y0)), cv2.FONT_HERSHEY_SIMPLEX, fontScale,
# color, thickness, cv2.LINE_AA)
# boxes = output[0]['boxes']
# scores = output[0]['scores']
# labels = output[0]['labels'].detach().cpu().numpy()
# for i, score in enumerate(scores):
# if score > 0.5:
# x0, y0, x1, y1 = boxes[i]
# cv2.rectangle(refine, (x0, y0), (x1, y1),
# (100, 200, 150), 1, 1)
# cv2.putText(refine, str(labels[i]), (x0, y0), cv2.FONT_HERSHEY_SIMPLEX, fontScale,u
# color, thickness, cv2.LINE_AA)
# combine = torch.concatenate(images[0],)
# writer.add_image('image',)
if j % 100 == 0:
print(
'Step {} -- loss_classifier = {} -- loss_box_reg = {} -- loss_objectness = {} -- loss_rpn_box_reg = {}\n'
.format(j, output['loss_classifier'].item(),
output['loss_box_reg'].item(),
output['loss_objectness'].item(),
output['loss_rpn_box_reg'].item()))
optimizer.zero_grad()
losses.backward()
optimizer.step()
step += 1
if i % 50 == 0:
torch.save(model.state_dict(), 'document.pth')
print('done')
gamma=0.1)
import os
# let's train it for 10 epochs
num_epochs = 10
train_losses = []
test_losses = []
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
#train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=1)
_, loss = train(model, optimizer, data_loader, device, epoch, print_freq=1)
train_losses.append(loss)
# # update the learning rate
lr_scheduler.step()
import matplotlib.pyplot as plt
# Plot training and val loss as a function of the epoch. Use this to monitor for overfitting.
plt.plot(range(1, num_epochs + 1), train_losses, label='traning')
plt.legend(loc='best')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss')
plt.savefig(os.path.join(root_path, 'prediction_visualization','loss.png'), bbox_inches='tight')
torch.save(model.state_dict(), "model_pollen_epoch{}.pt".format(num_epochs))
# test on testset
#evaluate(model, data_loader_test, device=device)