返回最大值的index pytorch方式

更新時間：2022年07月16日 15:37:45 作者：catbird233

這篇文章主要介紹了返回最大值的index pytorch方式，具有很好的參考價值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教

返回最大值的index

import torch
a=torch.tensor([[.1,.2,.3],
? ? ? ? ? ? ? ? [1.1,1.2,1.3],
? ? ? ? ? ? ? ? [2.1,2.2,2.3],
? ? ? ? ? ? ? ? [3.1,3.2,3.3]])
print(a.argmax(dim=1))
print(a.argmax())

輸出：

tensor([ 2, 2, 2, 2])
tensor(11)

pytorch 找最大值

題意：使用神經(jīng)網(wǎng)絡實現(xiàn)，從數(shù)組中找出最大值。

提供數(shù)據(jù)：兩個 csv 文件，一個存訓練集：n 個 m 維特征自然數(shù)數(shù)據(jù)，另一個存每條數(shù)據(jù)對應的 label ，就是每條數(shù)據(jù)中的最大值。

這里將隨機構建訓練集：

#%%
import numpy as np 
import pandas as pd 
import torch 
import random 
import torch.utils.data as Data
import torch.nn as nn
import torch.optim as optim
  
def GetData(m, n):
    dataset = []
    for j in range(m):
        max_v = random.randint(0, 9)
        data = [random.randint(0, 9) for i in range(n)]
        dataset.append(data)
    label = [max(dataset[i]) for i in  range(len(dataset))]
    data_list = np.column_stack((dataset, label))
    data_list = data_list.astype(np.float32)
    return data_list
 
#%%
# 數(shù)據(jù)集封裝 重載函數(shù)len, getitem
class GetMaxEle(Data.Dataset):
    def __init__(self, trainset):
        self.data = trainset 
 
    def __getitem__(self, index):
        item = self.data[index]
        x = item[:-1]
        y = item[-1]
        return x, y
    
    def __len__(self):
        return len(self.data)
 
# %% 定義網(wǎng)絡模型
class SingleNN(nn.Module):
    def __init__(self, n_feature, n_hidden, n_output):
        super(SingleNN, self).__init__()
        
        self.hidden = nn.Linear(n_feature, n_hidden)
        self.relu = nn.ReLU()
        self.predict = nn.Linear(n_hidden, n_output)
 
    def forward(self, x):
        x = self.hidden(x)
        x = self.relu(x)
        x = self.predict(x)
        return x
  
def train(m, n, batch_size, PATH):
    # 隨機生成 m 個 n 個維度的訓練樣本
    data_list =GetData(m, n)
    dataset = GetMaxEle(data_list)
    trainset = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                      shuffle=True)
 
    net = SingleNN(n_feature=10, n_hidden=100,
                   n_output=10)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
    #
    total_epoch = 100
    for epoch in range(total_epoch):
        for index, data in enumerate(trainset):
            input_x, labels = data
            labels = labels.long()
            optimizer.zero_grad()
 
            output = net(input_x)
            # print(output)
            # print(labels)
            loss = criterion(output, labels)
            loss.backward()
            optimizer.step()
 
        # scheduled_optimizer.step()
        print(f"Epoch {epoch}, loss:{loss.item()}")
 
    # %% 保存參數(shù)
    torch.save(net.state_dict(), PATH)
    #測試 
  
def test(m, n, batch_size, PATH):
    data_list = GetData(m, n)
    dataset = GetMaxEle(data_list)
    testloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
    dataiter = iter(testloader)
    input_x, labels = dataiter.next()
    net = SingleNN(n_feature=10, n_hidden=100,
                   n_output=10)
    net.load_state_dict(torch.load(PATH))
    outputs = net(input_x)
 
    _, predicted = torch.max(outputs, 1)
    print("Ground_truth:",labels.numpy())
    print("predicted:",predicted.numpy())
  
if __name__ == "__main__":
    m = 1000
    n = 10
    batch_size = 64
    PATH = './max_list.pth'
    train(m, n, batch_size, PATH)
    test(m, n, batch_size, PATH)

初始的想法是使用全連接網(wǎng)絡+分類來實現(xiàn)，但是結(jié)果不盡人意，主要原因：不同類別之間的樣本量差太大，幾乎90%都是最大值。

比如代碼中隨機構建 10 個 0~9 的數(shù)字構成一個樣本[2, 3, 5, 8, 9, 5, 3, 9, 3, 6], 該樣本標簽是9。

以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。

您可能感興趣的文章: