dataload/TensorLoader.lua at master · Element-Research/dataload · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
localdl=require'dataload._env'
localTensorLoader=torch.class('dl.TensorLoader', 'dl.DataLoader', dl)

functionTensorLoader:__init(inputs, targets)
self.inputs=inputs
self.targets=targets

assert(torchx.recursiveBatchSize(self.inputs) ==torchx.recursiveBatchSize(self.targets))
end

functionTensorLoader:index(indices, inputs, targets)
inputs=torchx.recursiveIndex(inputs, self.inputs, 1, indices)
targets=torchx.recursiveIndex(targets, self.targets, 1, indices)
returninputs, targets
end

functionTensorLoader:shuffle()
localindices=torch.LongTensor():randperm(self:size())
self.inputs=torchx.recursiveIndex(nil, self.inputs, 1, indices)
self.targets=torchx.recursiveIndex(nil, self.targets, 1, indices)
returnself, indices
end

functionTensorLoader:split(ratio)
assert(ratio>0andratio<1, "Expecting 0 < arg < 1")

localsize=self:size()
localsizeA=math.floor(size*ratio)

localloaders= {}
fori,splitinipairs{{1,sizeA},{sizeA+1,size}} do
localstart, stop=unpack(split)
localinputs=torchx.recursiveSub(self.inputs, start, stop)
localtargets=torchx.recursiveSub(self.targets, start, stop)
localloader=dl.TensorLoader(inputs, targets)
assert(loader:size() ==stop-start+1)
loaders[i] =loader
end
returnunpack(loaders)
end

functionTensorLoader:size()
returntorchx.recursiveBatchSize(self.inputs)
end

functionTensorLoader:isize(excludedim)
-- by default, batch dimension is excluded
excludedim=excludedim==niland1orexcludedim
returntorchx.recursiveSize(self.inputs, excludedim)
end

functionTensorLoader:tsize(excludedim)
-- by default, batch dimension is excluded
excludedim=excludedim==niland1orexcludedim
returntorchx.recursiveSize(self.targets, excludedim)
end