Focal loss的pytorch版本实现

最新推荐文章于 2024-08-19 23:40:59 发布

朴素.无恙

最新推荐文章于 2024-08-19 23:40:59 发布

阅读量3.9k

点赞数 2

分类专栏： pytorch 文章标签： Focal loss

本文链接：https://blog.csdn.net/weixin_40123108/article/details/91362628

版权

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision     
import numpy as np
import matplotlib.pyplot as plt 
import torch.nn.functional as F
torch.manual_seed(1)   
EPOCH = 3          
BATCH_SIZE = 50
LR = 0.001          
DOWNLOAD_MNIST = True 
train_data = torchvision.datasets.MNIST(
    root='./mnist/',   
    train=True, 
    transform=torchvision.transforms.ToTensor(),                                                     
    download=DOWNLOAD_MNIST,          
)
test_data = torchvision.datasets.MNIST(root='./mnist/', train=False)
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
test_x = Variable(torch.unsqueeze(test_data.test_data, dim=1), requires_grad=True).type(torch.FloatTensor)[:2000]/255 
test_y = test_data.test_labels[:2000]
class Net(nn.Module):
    #定义Net的初始化函数，这个函数定义了该神经网络的基本结构
    def __init__(self):
        super(Net, self).__init__() #复制并使用Net的父类的初始化方法，
                                    #即先运行nn.Module的初始化函数
        #self.conv1 = nn.Conv2d(1, 6, 5) # 定义conv1函数的是图像卷积函数：
                                        #输入为图像（1个频道，即灰度图）,输出为 6张特征图, 卷积核为5x5正方形
        #self.conv2 = nn.Conv2d(6, 16, 3)# 定义conv2函数的是图像卷积函数：
                                        #输入为6张特征图,输出为16张特征图, 卷积核为5x5正方形
        self.fc1   =